From b2710debaafd4046850ae1060ee0f9973ccf86b6 Mon Sep 17 00:00:00 2001
From: Ekaterina Tuzova <Ekaterina.Tuzova@jetbrains.com>
Date: Mon, 10 Jan 2011 16:29:24 +0300
Subject: [PATCH] fixed second part of PY-2648, and PY-2649

---
 .../psi/impl/PyAssignmentStatementImpl.java       | 12 +++++++++---
 python/testData/psi/assignment/TupleUnpack.py     |  6 +++++-
 .../jetbrains/python/PyAssignmentMappingTest.java | 15 +++++++++++----
 3 files changed, 25 insertions(+), 8 deletions(-)

diff --git a/python/src/com/jetbrains/python/psi/impl/PyAssignmentStatementImpl.java b/python/src/com/jetbrains/python/psi/impl/PyAssignmentStatementImpl.java
index a91cb53cc5a7..bdeed6b95aee 100644
--- a/python/src/com/jetbrains/python/psi/impl/PyAssignmentStatementImpl.java
+++ b/python/src/com/jetbrains/python/psi/impl/PyAssignmentStatementImpl.java
@@ -11,7 +11,6 @@ import com.intellij.util.SmartList;
 import com.jetbrains.python.PyTokenTypes;
 import com.jetbrains.python.psi.*;
 import com.jetbrains.python.toolbox.FP;
-import com.jetbrains.python.toolbox.RepeatIterable;
 import org.jetbrains.annotations.NotNull;
 import org.jetbrains.annotations.Nullable;
 
@@ -162,8 +161,15 @@ public class PyAssignmentStatementImpl extends PyElementImpl implements PyAssign
        map.add(new Pair<PyExpression, PyExpression>(lhs_one, rhs));
     }
     else if (lhs_tuple != null && rhs_one != null) { // multiple LHS, single RHS: unpacking
-      //for (PyExpression tuple_elt : lhs_tuple.getElements()) map.add(new Pair<PyExpression, PyExpression>(tuple_elt, rhs_one));
-      map.addAll(FP.zipList(Arrays.asList(lhs_tuple.getElements()), new RepeatIterable<PyExpression>(rhs_one)));
+      // PY-2648, PY-2649
+      PyElementGenerator elementGenerator = PyElementGenerator.getInstance(rhs_one.getProject());
+      int counter = 0;
+      for (PyExpression tuple_elt : lhs_tuple.getElements()) {
+        map.add(new Pair<PyExpression, PyExpression>(tuple_elt,
+                                                     elementGenerator.createExpressionFromText(rhs_one.getText() + "[" + counter + "]")));
+        ++counter;
+      }
+      //  map.addAll(FP.zipList(Arrays.asList(lhs_tuple.getElements()), new RepeatIterable<PyExpression>(rhs_one)));
     }
     else if (lhs_tuple != null && rhs_tuple != null) { // multiple both sides: piecewise mapping
       map.addAll(FP.zipList(Arrays.asList(lhs_tuple.getElements()), Arrays.asList(rhs_tuple.getElements()), null, null));
diff --git a/python/testData/psi/assignment/TupleUnpack.py b/python/testData/psi/assignment/TupleUnpack.py
index a0c96320be5a..366ec744052d 100644
--- a/python/testData/psi/assignment/TupleUnpack.py
+++ b/python/testData/psi/assignment/TupleUnpack.py
@@ -1 +1,5 @@
-<dst1>a, <dst2>b = <src>some_tuple
+def returnTuple():
+    return 5, 4
+
+<dst1>a, <dst2>b = <src>returnTuple()
+<src1>returnTuple()[0], <src2>returnTuple()[1]
\ No newline at end of file
diff --git a/python/testSrc/com/jetbrains/python/PyAssignmentMappingTest.java b/python/testSrc/com/jetbrains/python/PyAssignmentMappingTest.java
index 9020777e8d45..e4fd31e4bb51 100644
--- a/python/testSrc/com/jetbrains/python/PyAssignmentMappingTest.java
+++ b/python/testSrc/com/jetbrains/python/PyAssignmentMappingTest.java
@@ -122,7 +122,7 @@ public class PyAssignmentMappingTest extends LightMarkedTestCase {
       PsiElement dst = marks.get("<dst" + String.valueOf(i + 1) + ">").getParent(); // ident -> target expr
       assertTrue(dst instanceof PyTargetExpression);
       dsts[i] = dst;
-      PsiElement src = marks.get("<src" + String.valueOf(i + 1) +">").getParent(); // ident -> target expr
+      PsiElement src = marks.get("<src" + String.valueOf(i + 1) + ">").getParent(); // ident -> target expr
       assertTrue(src instanceof PyExpression);
       srcs[i] = src;
     }
@@ -161,21 +161,28 @@ public class PyAssignmentMappingTest extends LightMarkedTestCase {
   public void testTupleUnpack() throws Exception {
     Map<String, PsiElement> marks = loadTest();
     final int DST_NUM = 2;
-    assertEquals(DST_NUM+1, marks.size());
+    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" + String.valueOf(i+1) +">").getParent(); // ident -> target expr
       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" + String.valueOf(i+1) +">").getParent().getParent().getParent(); // ident -> target expr
+      assertTrue(src instanceof PyExpression);
+      srcs[i] = src;
+    }
+
     PsiElement src = marks.get("<src>").getParent(); // ident -> target expr
-    PyAssignmentStatement stmt = (PyAssignmentStatement)src.getParent();
+    PyAssignmentStatement stmt = (PyAssignmentStatement)src.getParent().getParent();
     List<Pair<PyExpression, PyExpression>> mapping = stmt.getTargetsToValuesMapping();
     assertEquals(DST_NUM, mapping.size());
     for (int i=0; i<DST_NUM; i+=1) {
       Pair<PyExpression, PyExpression> pair = mapping.get(i);
       assertEquals(dsts[i], pair.getFirst());
-      assertEquals(src, pair.getSecond());
+      assertEquals(srcs[i].getText(), pair.getSecond().getText());
     }
   }
 }