PY-75537 Implement PyAstAssignmentStatement.getTargets(), PyAstAssignmentStatement.getRawTargets() (PyFrontendElementTypesFacadeImpl.kt)

GitOrigin-RevId: 0f9292ac3ca0f4f4728b25cb2e5ddcd618b8319d
This commit is contained in:
Petr
2024-10-15 14:30:59 +02:00
committed by intellij-monorepo-bot
parent 85017a8482
commit 05942bee35
2 changed files with 62 additions and 60 deletions

View File

@@ -15,10 +15,13 @@
*/
package com.jetbrains.python.ast;
import com.intellij.lang.ASTNode;
import com.intellij.openapi.util.Pair;
import com.intellij.psi.PsiElement;
import com.intellij.psi.PsiErrorElement;
import com.intellij.psi.PsiNamedElement;
import com.intellij.psi.tree.TokenSet;
import com.jetbrains.python.PyTokenTypes;
import com.jetbrains.python.ast.impl.PyUtilCore;
import org.jetbrains.annotations.ApiStatus;
import org.jetbrains.annotations.NotNull;
@@ -35,6 +38,62 @@ import static com.jetbrains.python.ast.PyAstElementKt.findChildByClass;
@ApiStatus.Experimental
public interface PyAstAssignmentStatement extends PyAstStatement, PyAstNamedElementContainer, PyAstAnnotationOwner {
@ApiStatus.Internal
static <T extends PyAstExpression> T @NotNull [] calcTargets(@NotNull PyAstAssignmentStatement statement,
boolean raw,
T @NotNull [] result) {
final ASTNode[] eqSigns = statement.getNode().getChildren(TokenSet.create(PyTokenTypes.EQ));
if (eqSigns.length == 0) {
return result;
}
final ASTNode lastEq = eqSigns[eqSigns.length - 1];
List<PyAstExpression> candidates = new ArrayList<>();
ASTNode node = statement.getNode().getFirstChildNode();
while (node != null && node != lastEq) {
final PsiElement psi = node.getPsi();
if (psi instanceof PyAstExpression expression) {
if (raw) {
candidates.add(expression);
}
else {
addCandidate(candidates, expression);
}
}
node = node.getTreeNext();
}
List<PyAstExpression> targets = new ArrayList<>();
for (PyAstExpression expr : candidates) { // only filter out targets
if (raw ||
expr instanceof PyAstTargetExpression ||
expr instanceof PyAstReferenceExpression ||
expr instanceof PyAstSubscriptionExpression ||
expr instanceof PyAstSliceExpression) {
targets.add(expr);
}
}
return targets.toArray(result);
}
private static void addCandidate(List<PyAstExpression> candidates, PyAstExpression psi) {
if (psi instanceof PyAstParenthesizedExpression parenthesizedExpression) {
addCandidate(candidates, parenthesizedExpression.getContainedExpression());
}
else if (psi instanceof PyAstSequenceExpression sequenceExpression) {
for (PyAstExpression pyExpression : sequenceExpression.getElements()) {
addCandidate(candidates, pyExpression);
}
}
else if (psi instanceof PyAstStarExpression starExpression) {
final PyAstExpression expression = starExpression.getExpression();
if (expression != null) {
addCandidate(candidates, expression);
}
}
else {
candidates.add(psi);
}
}
@Nullable
@Override
default PyAstAnnotation getAnnotation() {