Fix nested lvalues and/or nested Unions

python · JukkaL · Oct 11, 2017 · Jul 24, 2017 · Oct 4, 2017 · Oct 5, 2017
commit a4b734ba8868b56da4c3bd9d3a384fe61ecb6275
diff --git a/mypy/binder.py b/mypy/binder.py
@@ -232,9 +232,12 @@ def accumulate_type_assignments(self) -> 'Iterator[Assigns]':
         the union are processed (a union of collected items is later bound
         manually by the caller).
         """
+        old_assignments = None
+        if self.type_assignments is not None:
+            old_assignments = self.type_assignments
         self.type_assignments = defaultdict(list)
         yield self.type_assignments
-        self.type_assignments = None
+        self.type_assignments = old_assignments
 
     def assign_type(self, expr: Expression,
                     type: Type,

diff --git a/mypy/checker.py b/mypy/checker.py
@@ -1639,7 +1639,7 @@ def check_multi_assignment(self, lvalues: List[Lvalue],
     def check_multi_assignment_from_union(self, lvalues: List[Expression], rvalue: Expression,
                                           rvalue_type: UnionType, context: Context,
                                           infer_lvalue_type: bool) -> None:
-        """Check assignment to multiple lavlue targets when rvalue type is a Union[...].
+        """Check assignment to multiple lvalue targets when rvalue type is a Union[...].
         For example:
 
             t: Union[Tuple[int, int], Tuple[str, str]]
@@ -1651,7 +1651,8 @@ def check_multi_assignment_from_union(self, lvalues: List[Expression], rvalue: E
         inferred types for first assignments, 'assignments' contains the narrowed types
         for binder.
         """
-        transposed = tuple([] for _ in lvalues)  # type: Tuple[List[Type], ...]
+        transposed = tuple([] for _ in
+                           self.flatten_lvalues(lvalues))  # type: Tuple[List[Type], ...]
         # Notify binder that we want to defer bindings and instead collect types.
         with self.binder.accumulate_type_assignments() as assignments:
             for item in rvalue_type.items:
@@ -1660,10 +1661,7 @@ def check_multi_assignment_from_union(self, lvalues: List[Expression], rvalue: E
                 self.check_multi_assignment(lvalues, rvalue, context,
                                             infer_lvalue_type=infer_lvalue_type,
                                             rv_type=item, undefined_rvalue=True)
-                for t, lv in zip(transposed, lvalues):
-                    if isinstance(lv, StarExpr):
-                        # Unwrap StarExpr, since it is unwrapped by other helpers.
-                        lv = lv.expr
+                for t, lv in zip(transposed, self.flatten_lvalues(lvalues)):
                     t.append(self.type_map.pop(lv, AnyType(TypeOfAny.special_form)))
         union_types = tuple(UnionType.make_simplified_union(col) for col in transposed)
         for expr, items in assignments.items():
@@ -1675,16 +1673,25 @@ def check_multi_assignment_from_union(self, lvalues: List[Expression], rvalue: E
                                     UnionType.make_simplified_union(types),
                                     UnionType.make_simplified_union(declared_types),
                                     False)
-        for union, lv in zip(union_ty
10000
pes, lvalues):
+        for union, lv in zip(union_types, self.flatten_lvalues(lvalues)):
             # Properly store the inferred types.
-            if isinstance(lv, StarExpr):
-                lv = lv.expr
             _1, _2, inferred = self.check_lvalue(lv)
             if inferred:
                 self.set_inferred_type(inferred, lv, union)
             else:
                 self.store_type(lv, union)
 
+    def flatten_lvalues(self, lvalues: List[Expression]) -> List[Expression]:
+        res = []  # type: List[Expression]
+        for lv in lvalues:
+            if isinstance(lv, (TupleExpr, ListExpr)):
+                res.extend(self.flatten_lvalues(lv.items))
+            if isinstance(lv, StarExpr):
+                # Unwrap StarExpr, since it is unwrapped by other helpers.
+                lv = lv.expr
+            res.append(lv)
+        return res
+
     def check_multi_assignment_from_tuple(self, lvalues: List[Lvalue], rvalue: Expression,
                                           rvalue_type: TupleType, context: Context,
                                           undefined_rvalue: bool,

diff --git a/test-data/unit/check-unions.test b/test-data/unit/check-unions.test
@@ -818,6 +818,49 @@ reveal_type(y) # E: Revealed type is 'builtins.object'
 [builtins fixtures/tuple.pyi]
 [out]
 
+[case testUnionUnpackingFromNestedTuples]
+from typing import Union, Tuple
+
+t: Union[Tuple[int, Tuple[int, int]], Tuple[str, Tuple[str, str]]]
+x, (y, z) = t
+reveal_type(x) # E: Revealed type is 'Union[builtins.int, builtins.str]'
+reveal_type(y) # E: Revealed type is 'Union[builtins.int, builtins.str]'
+reveal_type(z) # E: Revealed type is 'Union[builtins.int, builtins.str]'
+[builtins fixtures/tuple.pyi]
+[out]
+
+[case testNestedUnionUnpackingFromNestedTuples]
+from typing import Union, Tuple
+
+class A: pass
+class B: pass
+
+t: Union[Tuple[int, Union[Tuple[int, int], Tuple[A, A]]], Tuple[str, Union[Tuple[str, str], Tuple[B, B]]]]
+x, (y, z) = t
+reveal_type(x) # E: Revealed type is 'Union[builtins.int, builtins.str]'
+reveal_type(y) # E: Revealed type is 'Union[builtins.int, __main__.A, builtins.str, __main__.B]'
+reveal_type(z) # E: Revealed type is 'Union[builtins.int, __main__.A, builtins.str, __main__.B]'
+[builtins fixtures/tuple.pyi]
+[out]
+
+[case testNestedUnionUnpackingFromNestedTuplesBinder]
+from typing import Union, Tuple
+
+class A: pass
+class B: pass
+
+x: object
+y: object
+z: object
+
+t: Union[Tuple[int, Union[Tuple[int, int], Tuple[A, A]]], Tuple[str, Union[Tuple[str, str], Tuple[B, B]]]]
+x, (y, z) = t
+reveal_type(x) # E: Revealed type is 'Union[builtins.int, builtins.str]'
+reveal_type(y) # E: Revealed type is 'Union[builtins.int, __main__.A, builtins.str, __main__.B]'
+reveal_type(z) # E: Revealed type is 'Union[builtins.int, __main__.A, builtins.str, __main__.B]'
+[builtins fixtures/tuple.pyi]
+[out]
+
 [case testLongUnionFormatting]
 from typing import Any, Generic, TypeVar, Union