[python] PyAssignmentStatementImpl.getTargetsToValuesMapping() Unpack nested tuples as specified in javadoc

GitOrigin-RevId: e0d4c13cadfa11d85ba3ddeba17e3c4cdda0a46d

python/python-psi-impl/src/com/jetbrains/python/psi/impl/PyAssignmentStatementImpl.java

@@ -13,8 +13,6 @@ import com.intellij.util.SmartList;
 import com.intellij.util.containers.ContainerUtil;
 import com.jetbrains.python.PyTokenTypes;
 import com.jetbrains.python.psi.*;
-import one.util.streamex.IntStreamEx;
-import one.util.streamex.StreamEx;
 import org.jetbrains.annotations.NotNull;
 import org.jetbrains.annotations.Nullable;
 
@@ -119,17 +117,17 @@ public class PyAssignmentStatementImpl extends PyElementImpl implements PyAssign
     return ret;
   }
 
-  private static void mapToValues(PyExpression lhs, PyExpression rhs, List<Pair<PyExpression, PyExpression>> map) {
+  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 PySequenceExpression tupleExpr) lhs_tuple = tupleExpr;
+    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 PySequenceExpression tupleExpr) rhs_tuple = tupleExpr;
+    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)
@@ -144,7 +142,7 @@ public class PyAssignmentStatementImpl extends PyElementImpl implements PyAssign
         try {
           final PyExpression expression =
             elementGenerator.createExpressionFromText(languageLevel, "(" + rhs_one.getText() + ")[" + counter + "]");
-          map.add(Pair.create(tuple_elt, expression));
+          mapToValues(tuple_elt, expression, map);
         }
         catch (IncorrectOperationException e) {
           // not parsed, no problem
@@ -156,18 +154,13 @@ public class PyAssignmentStatementImpl extends PyElementImpl implements PyAssign
       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());
-
-      map.addAll(StreamEx.zip(alignToSize(lhsTupleElements, size), alignToSize(rhsTupleElements, size), Pair::create).toList());
+      for (int index = 0; index < size; index++) {
+        mapToValues(ContainerUtil.getOrElse(lhsTupleElements, index, null),
+                    ContainerUtil.getOrElse(rhsTupleElements, index, null), map);
+      }
     }
   }
 
-  @NotNull
-  private static <T> List<T> alignToSize(@NotNull List<T> list, int size) {
-    return list.size() == size
-           ? list
-           : IntStreamEx.range(size).mapToObj(index -> ContainerUtil.getOrElse(list, index, null)).toList();
-  }
-
   @Nullable
   public PsiNamedElement getNamedElement(@NotNull final String the_name) {
     // performance: check simple case first

python/testData/psi/assignment/NestedTupleMapped.py

@@ -0,0 +1 @@
+<dst1>a, (<dst2>b, <dst3>c) = <src1>1, (<src2>2, <src3>3)

python/testData/psi/assignment/NestedTupleUnpack.py

@@ -0,0 +1,5 @@
+def returnTuple():
+    pass
+
+<dst1>a, (<dst2>b, (<dst3>c, <dst4>d)) = <src1>1, <src>returnTuple()
+<src2>((returnTuple())[0]), <src3>(((returnTuple())[1])[0]), <src4>(((returnTuple())[1])[1])

python/testData/psi/assignment/TupleUnpack.py

@@ -2,4 +2,4 @@ def returnTuple():
     return 5, 4
 
 <dst1>a, <dst2>b = <src>returnTuple()
-<src1>(returnTuple())[0], <src2>(returnTuple())[1]
+<src1>((returnTuple())[0]), <src2>((returnTuple())[1])

python/testSrc/com/jetbrains/python/PyAssignmentMappingTest.java

@@ -18,12 +18,18 @@ package com.jetbrains.python;
 import com.intellij.openapi.util.Pair;
 import com.intellij.psi.PsiElement;
 import com.intellij.psi.util.PsiTreeUtil;
+import com.intellij.util.containers.ContainerUtil;
 import com.jetbrains.python.fixtures.LightMarkedTestCase;
-import com.jetbrains.python.psi.*;
+import com.jetbrains.python.psi.PyAssignmentStatement;
+import com.jetbrains.python.psi.PyExpression;
+import com.jetbrains.python.psi.PySubscriptionExpression;
+import com.jetbrains.python.psi.PyTargetExpression;
+import com.jetbrains.python.psi.impl.PyPsiUtils;
 import junit.framework.Assert;
 
 import java.util.List;
 import java.util.Map;
+import java.util.stream.IntStream;
 
 /**
  * Tests assignment mapping.
@@ -101,51 +107,23 @@ public class PyAssignmentMappingTest extends LightMarkedTestCase {
   }
 
   public void testTupleMapped() {
-    Map<String, PsiElement> marks = loadTest();
-    final int PAIR_NUM = 2;
-    Assert.assertEquals(PAIR_NUM * 2, marks.size());
-    PsiElement[] srcs = new PsiElement[PAIR_NUM];
-    PsiElement[] dsts = new PsiElement[PAIR_NUM];
-    for (int i=0; i<PAIR_NUM; i+=1) {
-      PsiElement dst = marks.get("<dst" + (i + 1) + ">").getParent(); // ident -> target expr
-      Assert.assertTrue(dst instanceof PyTargetExpression);
-      dsts[i] = dst;
-      PsiElement src = marks.get("<src" + (i + 1) + ">").getParent(); // ident -> target expr
-      Assert.assertTrue(src instanceof PyExpression);
-      srcs[i] = src;
-    }
-    PyAssignmentStatement stmt = (PyAssignmentStatement)srcs[0].getParent().getParent(); // tuple expr -> assignment
-    List<Pair<PyExpression, PyExpression>> mapping = stmt.getTargetsToValuesMapping();
-    Assert.assertEquals(PAIR_NUM, mapping.size());
-    for (int i=0; i<PAIR_NUM; i+=1) {
-      Pair<PyExpression, PyExpression> pair = mapping.get(i);
-      Assert.assertEquals(dsts[i], pair.getFirst());
-      Assert.assertEquals(srcs[i], pair.getSecond());
-    }
+    List<Pair<PyExpression, PyExpression>> expectedMappings = loadMultiMappingTest(IntStream.of(1, 2));
+    PyAssignmentStatement stmt =
+      PsiTreeUtil.getParentOfType(expectedMappings.get(0).second, PyAssignmentStatement.class);
+    assertSameElements(stmt.getTargetsToValuesMapping(), expectedMappings);
+  }
+
+  public void testNestedTupleMapped() {
+    List<Pair<PyExpression, PyExpression>> expectedMappings = loadMultiMappingTest(IntStream.rangeClosed(1, 3));
+    PyAssignmentStatement stmt =
+      PsiTreeUtil.getParentOfType(expectedMappings.get(0).second, PyAssignmentStatement.class);
+    assertSameElements(stmt.getTargetsToValuesMapping(), expectedMappings);
   }
 
   public void testParenthesizedTuple() { //PY-2648
-    Map<String, PsiElement> marks = loadTest();
-    final int PAIR_NUM = 2;
-    Assert.assertEquals(PAIR_NUM * 2, marks.size());
-    PsiElement[] srcs = new PsiElement[PAIR_NUM];
-    PsiElement[] dsts = new PsiElement[PAIR_NUM];
-    for (int i=0; i<PAIR_NUM; i+=1) {
-      PsiElement dst = marks.get("<dst" + (i + 1) + ">").getParent(); // ident -> target expr
-      Assert.assertTrue(dst instanceof PyTargetExpression);
-      dsts[i] = dst;
-      PsiElement src = marks.get("<src" + (i + 1) + ">").getParent(); // ident -> target expr
-      Assert.assertTrue(src instanceof PyExpression);
-      srcs[i] = src;
-    }
-    PyAssignmentStatement stmt = PsiTreeUtil.getParentOfType(srcs[0], PyAssignmentStatement.class);
-    List<Pair<PyExpression, PyExpression>> mapping = stmt.getTargetsToValuesMapping();
-    Assert.assertEquals(PAIR_NUM, mapping.size());
-    for (int i=0; i<PAIR_NUM; i+=1) {
-      Pair<PyExpression, PyExpression> pair = mapping.get(i);
-      Assert.assertEquals(dsts[i], pair.getFirst());
-      Assert.assertEquals(srcs[i], pair.getSecond());
-    }
+    List<Pair<PyExpression, PyExpression>> expectedMappings = loadMultiMappingTest(IntStream.of(1, 2));
+    PyAssignmentStatement stmt = PsiTreeUtil.getParentOfType(expectedMappings.get(0).second, PyAssignmentStatement.class);
+    assertSameElements(stmt.getTargetsToValuesMapping(), expectedMappings);
   }
 
   public void testTuplePack() {
@@ -169,32 +147,46 @@ public class PyAssignmentMappingTest extends LightMarkedTestCase {
     }
   }
 
-
   public void testTupleUnpack() {
     Map<String, PsiElement> marks = loadTest();
-    final int DST_NUM = 2;
-    Assert.assertEquals(DST_NUM + 3, marks.size());
-    PsiElement[] dsts = new PsiElement[DST_NUM];
-    for (int i=0; i<DST_NUM; i+=1) {
-      PsiElement dst = marks.get("<dst" + (i + 1) + ">").getParent(); // ident -> target expr
-      Assert.assertTrue(dst instanceof PyTargetExpression);
-      dsts[i] = dst;
-    }
-    PsiElement[] srcs = new PsiElement[DST_NUM];
-    for (int i=0; i<DST_NUM; i+=1) {
-      PsiElement src = marks.get("<src" + (i + 1) + ">").getParent().getParent(); // ident -> target expr
-      Assert.assertTrue(src instanceof PyExpression);
-      srcs[i] = src;
-    }
+    List<Pair<PyExpression, PyExpression>> expectedMapping = getMapping(marks, IntStream.rangeClosed(1, 2));
 
-    PsiElement src = marks.get("<src>").getParent(); // ident -> target expr
+    PsiElement src = marks.get("<src>").getParent();
     PyAssignmentStatement stmt = (PyAssignmentStatement)src.getParent().getParent();
     List<Pair<PyExpression, PyExpression>> mapping = stmt.getTargetsToValuesMapping();
-    Assert.assertEquals(DST_NUM, mapping.size());
-    for (int i=0; i<DST_NUM; i+=1) {
-      Pair<PyExpression, PyExpression> pair = mapping.get(i);
-      Assert.assertEquals(dsts[i], pair.getFirst());
-      Assert.assertEquals(srcs[i].getText(), pair.getSecond().getText());
-    }
+
+    assertSameElements(
+      ContainerUtil.map(mapping, pair -> Pair.create(pair.first, pair.second.getText())),
+      ContainerUtil.map(expectedMapping, pair -> Pair.create(pair.first, pair.second.getText()))
+    );
+  }
+
+  public void testNestedTupleUnpack() {
+    Map<String, PsiElement> marks = loadTest();
+    List<Pair<PyExpression, PyExpression>> expectedMapping = getMapping(marks, IntStream.rangeClosed(1, 4));
+
+    PsiElement src = marks.get("<src>").getParent();
+    PyAssignmentStatement stmt = PsiTreeUtil.getParentOfType(src, PyAssignmentStatement.class);
+    List<Pair<PyExpression, PyExpression>> mapping = stmt.getTargetsToValuesMapping();
+
+    assertSameElements(
+      ContainerUtil.map(mapping, pair -> Pair.create(pair.first, pair.second.getText())),
+      ContainerUtil.map(expectedMapping, pair -> Pair.create(pair.first, pair.second.getText()))
+    );
+  }
+
+  private List<Pair<PyExpression, PyExpression>> loadMultiMappingTest(IntStream indices) {
+    Map<String, PsiElement> marks = loadTest();
+    return getMapping(marks, indices);
+  }
+
+  private static List<Pair<PyExpression, PyExpression>> getMapping(Map<String, PsiElement> marks, IntStream indices) {
+    return indices.mapToObj(i -> Pair.create(getMarkedExpression(marks, "<dst%d>".formatted(i)),
+                                             getMarkedExpression(marks, "<src%d>".formatted(i)))).toList();
+  }
+
+  private static PyExpression getMarkedExpression(Map<String, PsiElement> marks, String marker) {
+    PsiElement element = marks.get(marker);
+    return PyPsiUtils.flattenParens((PyExpression)element.getParent());
   }
 }