PY-75537 Implement PyAstAssignmentStatement.getTargetsToValuesMapping() (PyFrontendElementTypesFacadeImpl.kt)

GitOrigin-RevId: 8c4926689e42f117275986c2e1246d7e04bde489
This commit is contained in:
Petr
2024-10-15 18:26:09 +02:00
committed by intellij-monorepo-bot
parent 47e2c8c0bc
commit 5befce7fb9
6 changed files with 89 additions and 75 deletions

View File

@@ -21,13 +21,21 @@ import com.intellij.psi.PsiElement;
import com.intellij.psi.PsiErrorElement;
import com.intellij.psi.PsiNamedElement;
import com.intellij.psi.tree.TokenSet;
import com.intellij.psi.util.PsiTreeUtil;
import com.intellij.util.IncorrectOperationException;
import com.intellij.util.SmartList;
import com.intellij.util.containers.ContainerUtil;
import com.jetbrains.python.PyTokenTypes;
import com.jetbrains.python.ast.impl.PyPsiUtilsCore;
import com.jetbrains.python.ast.impl.PyUtilCore;
import com.jetbrains.python.psi.LanguageLevel;
import com.jetbrains.python.psi.PyAstElementGenerator;
import org.jetbrains.annotations.ApiStatus;
import org.jetbrains.annotations.NotNull;
import org.jetbrains.annotations.Nullable;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.List;
import static com.jetbrains.python.ast.PyAstElementKt.findChildByClass;
@@ -141,7 +149,66 @@ public interface PyAstAssignmentStatement extends PyAstStatement, PyAstNamedElem
* @return a list of [target, value] pairs; either part of a pair may be null, but not both.
*/
@NotNull
List<? extends Pair<? extends PyAstExpression, ? extends PyAstExpression>> getTargetsToValuesMapping();
default List<? extends Pair<? extends PyAstExpression, ? extends PyAstExpression>> getTargetsToValuesMapping() {
List<Pair<PyAstExpression, PyAstExpression>> ret = new SmartList<>();
if (!PsiTreeUtil.hasErrorElements(this)) { // no parse errors
PyAstExpression[] constituents = PsiTreeUtil.getChildrenOfType(this, PyAstExpression.class); // "a = b = c" -> [a, b, c]
if (constituents != null && constituents.length > 1) {
int lastIndex = constituents.length - 1;
PyAstExpression rhs = constituents[lastIndex];
for (int i = 0; i < lastIndex; i++) {
mapToValues(constituents[i], rhs, ret);
}
}
}
return ret;
}
private static void mapToValues(@Nullable PyAstExpression lhs,
@Nullable PyAstExpression rhs,
List<Pair<PyAstExpression, PyAstExpression>> map) {
// cast for convenience
PyAstSequenceExpression lhs_tuple = null;
PyAstExpression lhs_one = null;
if (PyPsiUtilsCore.flattenParens(lhs) instanceof PyAstTupleExpression<?> tupleExpr) lhs_tuple = tupleExpr;
else if (lhs != null) lhs_one = lhs;
PyAstSequenceExpression rhs_tuple = null;
PyAstExpression rhs_one = null;
if (PyPsiUtilsCore.flattenParens(rhs) instanceof PyAstTupleExpression<?> tupleExpr) rhs_tuple = tupleExpr;
else if (rhs != null) rhs_one = rhs;
//
if (lhs_one != null) { // single LHS, single RHS (direct mapping) or multiple RHS (packing)
map.add(Pair.create(lhs_one, rhs));
}
else if (lhs_tuple != null && rhs_one != null) { // multiple LHS, single RHS: unpacking
// PY-2648, PY-2649
PyAstElementGenerator elementGenerator = PyAstElementGenerator.getInstance(rhs_one.getProject());
final LanguageLevel languageLevel = LanguageLevel.forElement(lhs);
int counter = 0;
for (PyAstExpression tuple_elt : lhs_tuple.getElements()) {
try {
final PyAstExpression expression =
elementGenerator.createExpressionFromText(languageLevel, "(" + rhs_one.getText() + ")[" + counter + "]");
mapToValues(tuple_elt, expression, map);
}
catch (IncorrectOperationException e) {
// not parsed, no problem
}
++counter;
}
}
else if (lhs_tuple != null && rhs_tuple != null) { // multiple both sides: piecewise mapping
final List<PyAstExpression> lhsTupleElements = Arrays.asList(lhs_tuple.getElements());
final List<PyAstExpression> rhsTupleElements = Arrays.asList(rhs_tuple.getElements());
final int size = Math.max(lhsTupleElements.size(), rhsTupleElements.size());
for (int index = 0; index < size; index++) {
mapToValues(ContainerUtil.getOrElse(lhsTupleElements, index, null),
ContainerUtil.getOrElse(rhsTupleElements, index, null), map);
}
}
}
@Nullable
default PyAstExpression getLeftHandSideExpression() {

View File

@@ -8,8 +8,10 @@ import com.intellij.psi.PsiFile;
import com.intellij.psi.PsiFileFactory;
import com.intellij.psi.impl.PsiFileFactoryImpl;
import com.intellij.testFramework.LightVirtualFile;
import com.intellij.util.IncorrectOperationException;
import com.jetbrains.python.PythonFileType;
import com.jetbrains.python.PythonLanguage;
import com.jetbrains.python.ast.PyAstExpression;
import com.jetbrains.python.ast.PyAstExpressionStatement;
import org.jetbrains.annotations.ApiStatus;
import org.jetbrains.annotations.NotNull;
@@ -126,4 +128,15 @@ public class PyAstElementGenerator {
}
protected void specifyFileLanguageLevel(@NotNull VirtualFile virtualFile, @Nullable LanguageLevel langLevel) { }
@NotNull
public PyAstExpression createExpressionFromText(@NotNull LanguageLevel languageLevel, @NotNull String text)
throws IncorrectOperationException {
final PsiFile dummyFile = createDummyFile(languageLevel, text);
final PsiElement element = dummyFile.getFirstChild();
if (element instanceof PyAstExpressionStatement expressionStatement) {
return expressionStatement.getExpression();
}
throw new IncorrectOperationException("could not parse text as expression: " + text);
}
}

View File

@@ -84,7 +84,10 @@ public interface PyAssignmentStatement extends PyAstAssignmentStatement, PyState
*/
@Override
@NotNull
List<Pair<PyExpression, PyExpression>> getTargetsToValuesMapping();
default List<Pair<PyExpression, PyExpression>> getTargetsToValuesMapping() {
//noinspection unchecked
return (List<Pair<PyExpression, PyExpression>>)PyAstAssignmentStatement.super.getTargetsToValuesMapping();
}
@Override
@Nullable

View File

@@ -52,8 +52,11 @@ public abstract class PyElementGenerator extends PyAstElementGenerator {
public abstract PyBinaryExpression createBinaryExpression(String s, PyExpression expr, PyExpression listLiteral);
@Override
@NotNull
public abstract PyExpression createExpressionFromText(@NotNull LanguageLevel languageLevel, @NotNull String text) throws IncorrectOperationException;
public PyExpression createExpressionFromText(@NotNull LanguageLevel languageLevel, @NotNull String text) throws IncorrectOperationException {
return (PyExpression)super.createExpressionFromText(languageLevel, text);
}
@NotNull
public abstract PyPattern createPatternFromText(@NotNull LanguageLevel languageLevel, @NotNull String text) throws IncorrectOperationException;

View File

@@ -44,67 +44,6 @@ public class PyAssignmentStatementImpl extends PyElementImpl implements PyAssign
return PyAstAssignmentStatement.calcTargets(this, true, PyExpression.EMPTY_ARRAY);
}
@Override
@NotNull
public List<Pair<PyExpression, PyExpression>> getTargetsToValuesMapping() {
List<Pair<PyExpression, PyExpression>> ret = new SmartList<>();
if (!PsiTreeUtil.hasErrorElements(this)) { // no parse errors
PyExpression[] constituents = PsiTreeUtil.getChildrenOfType(this, PyExpression.class); // "a = b = c" -> [a, b, c]
if (constituents != null && constituents.length > 1) {
int lastIndex = constituents.length - 1;
PyExpression rhs = constituents[lastIndex];
for (int i = 0; i < lastIndex; i++) {
mapToValues(constituents[i], rhs, ret);
}
}
}
return ret;
}
private static void mapToValues(@Nullable PyExpression lhs, @Nullable PyExpression rhs, List<Pair<PyExpression, PyExpression>> map) {
// cast for convenience
PySequenceExpression lhs_tuple = null;
PyExpression lhs_one = null;
if (PyPsiUtils.flattenParens(lhs) instanceof PyTupleExpression tupleExpr) lhs_tuple = tupleExpr;
else if (lhs != null) lhs_one = lhs;
PySequenceExpression rhs_tuple = null;
PyExpression rhs_one = null;
if (PyPsiUtils.flattenParens(rhs) instanceof PyTupleExpression tupleExpr) rhs_tuple = tupleExpr;
else if (rhs != null) rhs_one = rhs;
//
if (lhs_one != null) { // single LHS, single RHS (direct mapping) or multiple RHS (packing)
map.add(Pair.create(lhs_one, rhs));
}
else if (lhs_tuple != null && rhs_one != null) { // multiple LHS, single RHS: unpacking
// PY-2648, PY-2649
PyElementGenerator elementGenerator = PyElementGenerator.getInstance(rhs_one.getProject());
final LanguageLevel languageLevel = LanguageLevel.forElement(lhs);
int counter = 0;
for (PyExpression tuple_elt : lhs_tuple.getElements()) {
try {
final PyExpression expression =
elementGenerator.createExpressionFromText(languageLevel, "(" + rhs_one.getText() + ")[" + counter + "]");
mapToValues(tuple_elt, expression, map);
}
catch (IncorrectOperationException e) {
// not parsed, no problem
}
++counter;
}
}
else if (lhs_tuple != null && rhs_tuple != null) { // multiple both sides: piecewise mapping
final List<PyExpression> lhsTupleElements = Arrays.asList(lhs_tuple.getElements());
final List<PyExpression> rhsTupleElements = Arrays.asList(rhs_tuple.getElements());
final int size = Math.max(lhsTupleElements.size(), rhsTupleElements.size());
for (int index = 0; index < size; index++) {
mapToValues(ContainerUtil.getOrElse(lhsTupleElements, index, null),
ContainerUtil.getOrElse(rhsTupleElements, index, null), map);
}
}
}
@Nullable
public PsiNamedElement getNamedElement(@NotNull final String the_name) {
// performance: check simple case first

View File

@@ -234,17 +234,6 @@ public final class PyElementGeneratorImpl extends PyElementGenerator {
return binExpr;
}
@Override
@NotNull
public PyExpression createExpressionFromText(@NotNull LanguageLevel languageLevel, @NotNull String text) {
final PsiFile dummyFile = createDummyFile(languageLevel, text);
final PsiElement element = dummyFile.getFirstChild();
if (element instanceof PyExpressionStatement) {
return ((PyExpressionStatement)element).getExpression();
}
throw new IncorrectOperationException("could not parse text as expression: " + text);
}
@Override
public @NotNull PyPattern createPatternFromText(@NotNull LanguageLevel languageLevel, @NotNull String text)
throws IncorrectOperationException {