From a1d93e0756574c4e5913ffdfc34d1d4443c06b69 Mon Sep 17 00:00:00 2001
From: Andrey Vlasovskikh <andrey.vlasovskikh@jetbrains.com>
Date: Tue, 27 May 2014 19:32:26 +0400
Subject: [PATCH] Use union of types of members as iteration type for tuples

---
 .../jetbrains/python/psi/impl/PyTargetExpressionImpl.java | 8 ++++++++
 python/testSrc/com/jetbrains/python/PyTypeTest.java       | 7 +++++++
 2 files changed, 15 insertions(+)

diff --git a/python/src/com/jetbrains/python/psi/impl/PyTargetExpressionImpl.java b/python/src/com/jetbrains/python/psi/impl/PyTargetExpressionImpl.java
index 246d168ea8cd..545e8e86cba1 100644
--- a/python/src/com/jetbrains/python/psi/impl/PyTargetExpressionImpl.java
+++ b/python/src/com/jetbrains/python/psi/impl/PyTargetExpressionImpl.java
@@ -386,6 +386,14 @@ public class PyTargetExpressionImpl extends PyPresentableElementImpl<PyTargetExp
         }
       }
     }
+    else if (iterableType instanceof PyTupleType) {
+      final PyTupleType tupleType = (PyTupleType)iterableType;
+      final List<PyType> memberTypes = new ArrayList<PyType>();
+      for (int i = 0; i < tupleType.getElementCount(); i++) {
+        memberTypes.add(tupleType.getElementType(i));
+      }
+      return PyUnionType.union(memberTypes);
+    }
     else if (iterableType instanceof PyUnionType) {
       final Collection<PyType> members = ((PyUnionType)iterableType).getMembers();
       final List<PyType> iterationTypes = new ArrayList<PyType>();
diff --git a/python/testSrc/com/jetbrains/python/PyTypeTest.java b/python/testSrc/com/jetbrains/python/PyTypeTest.java
index e82ecaa81786..81b4aa46629e 100644
--- a/python/testSrc/com/jetbrains/python/PyTypeTest.java
+++ b/python/testSrc/com/jetbrains/python/PyTypeTest.java
@@ -812,6 +812,13 @@ public class PyTypeTest extends PyTestCase {
            "  expr, foo = xs\n");
   }
 
+  public void testTupleIterationType() {
+    doTest("int | str",
+           "xs = (1, 'a')\n" +
+           "for expr in xs:\n" +
+           "    pass\n");
+  }
+
   private static TypeEvalContext getTypeEvalContext(@NotNull PyExpression element) {
     return TypeEvalContext.userInitiated(element.getContainingFile()).withTracing();
   }