Allow assignments to multiple targets from union types (#4067)

Fixes #3859
Fixes #3240
Fixes #1855
Fixes #1575

This is a simple fix of various bugs and a crash based on the idea
originally proposed in #2219. The idea is to check assignment for 
every item in a union. However, in contrast to #2219, I think it will 
be more expected/consistent to construct a union of the resulting 
types instead of a join, for example:

```
x: Union[int, str]
x1 = x
reveal_type(x1) # Revealed type is 'Union[int, str]'

y: Union[Tuple[int], Tuple[str]]
(y1,) = y
reveal_type(y1) # Revealed type is 'Union[int, str]'
```

@elazarg did the initial work on this.

Author: ilevkivskyi

mypy/binder.py

@@ -1,5 +1,10 @@
-from typing import Dict, List, Set, Iterator, Union, Optional, cast
+from typing import Dict, List, Set, Iterator, Union, Optional, Tuple, cast
 from contextlib import contextmanager
+from collections import defaultdict
+
+MYPY = False
+if MYPY:
+    from typing import DefaultDict
 
 from mypy.types import Type, AnyType, PartialType, UnionType, TypeOfAny
 from mypy.subtypes import is_subtype
@@ -37,6 +42,12 @@ def __init__(self) -> None:
         self.unreachable = False
 
 
+if MYPY:
+    # This is the type of stored assignments for union type rvalues.
+    # We use 'if MYPY: ...' since typing-3.5.1 does not have 'DefaultDict'
+    Assigns = DefaultDict[Expression, List[Tuple[Type, Optional[Type]]]]
+
+
 class ConditionalTypeBinder:
     """Keep track of conditional types of variables.
 
@@ -57,6 +68,9 @@ class A:
     reveal_type(lst[0].a) # str
     ```
     """
+    # Stored assignments for situations with tuple/list lvalue and rvalue of union type.
+    # This maps an expression to a list of bound types for every item in the union type.
+    type_assignments = None  # type: Optional[Assigns]
 
     def __init__(self) -> None:
         # The stack of frames currently used.  These map
@@ -210,10 +224,30 @@ def pop_frame(self, can_skip: bool, fall_through: int) -> Frame:
 
         return result
 
+    @contextmanager
+    def accumulate_type_assignments(self) -> 'Iterator[Assigns]':
+        """Push a new map to collect assigned types in multiassign from union.
+
+        If this map is not None, actual binding is deferred until all items in
+        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 = old_assignments
+
     def assign_type(self, expr: Expression,
                     type: Type,
                     declared_type: Optional[Type],
                     restrict_any: bool = False) -> None:
+        if self.type_assignments is not None:
+            # We are in a multiassign from union, defer the actual binding,
+            # just collect the types.
+            self.type_assignments[expr].append((type, declared_type))
+            return
         if not isinstance(expr, BindableTypes):
             return None
         if not literal(expr):

mypy/checker.py

@@ -131,6 +131,8 @@ class TypeChecker(NodeVisitor[None], CheckerPluginInterface):
     # Used for collecting inferred attribute types so that they can be checked
     # for consistency.
     inferred_attribute_types = None  # type: Optional[Dict[Var, Type]]
+    # Don't infer partial None types if we are processing assignment from Union
+    no_partial_types = False  # type: bool
 
     # The set of all dependencies (suppressed or not) that this module accesses, either
     # directly or indirectly.
@@ -1605,12 +1607,13 @@ def check_multi_assignment(self, lvalues: List[Lvalue],
                                rvalue: Expression,
                                context: Context,
                                infer_lvalue_type: bool = True,
-                               msg: Optional[str] = None) -> None:
+                               rv_type: Optional[Type] = None,
+                               undefined_rvalue: bool = False) -> None:
         """Check the assignment of one rvalue to a number of lvalues."""
 
         # Infer the type of an ordinary rvalue expression.
-        rvalue_type = self.expr_checker.accept(rvalue)  # TODO maybe elsewhere; redundant
-        undefined_rvalue = False
+        # TODO: maybe elsewhere; redundant.
+        rvalue_type = rv_type or self.expr_checker.accept(rvalue)
 
         if isinstance(rvalue_type, UnionType):
             # If this is an Optional type in non-strict Optional code, unwrap it.
@@ -1628,10 +1631,71 @@ def check_multi_assignment(self, lvalues: List[Lvalue],
         elif isinstance(rvalue_type, TupleType):
             self.check_multi_assignment_from_tuple(lvalues, rvalue, rvalue_type,
                                                    context, undefined_rvalue, infer_lvalue_type)
+        elif isinstance(rvalue_type, UnionType):
+            self.check_multi_assignment_from_union(lvalues, rvalue, rvalue_type, context,
+                                                   infer_lvalue_type)
         else:
             self.check_multi_assignment_from_iterable(lvalues, rvalue_type,
                                                       context, infer_lvalue_type)
 
+    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 lvalue targets when rvalue type is a Union[...].
+        For example:
+
+            t: Union[Tuple[int, int], Tuple[str, str]]
+            x, y = t
+            reveal_type(x)  # Union[int, str]
+
+        The idea in this case is to process the assignment for every item of the union.
+        Important note: the types are collected in two places, 'union_types' contains
+        inferred types for first assignments, 'assignments' contains the narrowed types
+        for binder.
+        """
+        self.no_partial_types = True
+        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:
+                # Type check the assignment separately for each union item and collect
+                # the inferred lvalue types for each union item.
+                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, 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():
+            # Bind a union of types collected in 'assignments' to every expression.
+            if isinstance(expr, StarExpr):
+                expr = expr.expr
+            types, declared_types = zip(*items)
+            self.binder.assign_type(expr,
+                                    UnionType.make_simplified_union(types),
+                                    UnionType.make_simplified_union(declared_types),
+                                    False)
+        for union, lv in zip(union_types, self.flatten_lvalues(lvalues)):
+            # Properly store the inferred types.
+            _1, _2, inferred = self.check_lvalue(lv)
+            if inferred:
+                self.set_inferred_type(inferred, lv, union)
+            else:
+                self.store_type(lv, union)
+        self.no_partial_types = False
+
+    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,
@@ -1654,7 +1718,11 @@ def check_multi_assignment_from_tuple(self, lvalues: List[Lvalue], rvalue: Expre
                     relevant_items = reinferred_rvalue_type.relevant_items()
                     if len(relevant_items) == 1:
                         reinferred_rvalue_type = relevant_items[0]
-
+                if isinstance(reinferred_rvalue_type, UnionType):
+                    self.check_multi_assignment_from_union(lvalues, rvalue,
+                                                           reinferred_rvalue_type, context,
+                                                           infer_lvalue_type)
+                    return
                 assert isinstance(reinferred_rvalue_type, TupleType)
                 rvalue_type = reinferred_rvalue_type
 
@@ -1716,7 +1784,7 @@ def split_around_star(self, items: List[T], star_index: int,
         returns in: ([1,2], [3,4,5], [6,7])
         """
         nr_right_of_star = length - star_index - 1
-        right_index = nr_right_of_star if -nr_right_of_star != 0 else len(items)
+        right_index = -nr_right_of_star if nr_right_of_star != 0 else len(items)
         left = items[:star_index]
         star = items[star_index:right_index]
         right = items[right_index:]
@@ -1800,7 +1868,7 @@ def infer_variable_type(self, name: Var, lvalue: Lvalue,
         """Infer the type of initialized variables from initializer type."""
         if isinstance(init_type, DeletedType):
             self.msg.deleted_as_rvalue(init_type, context)
-        elif not is_valid_inferred_type(init_type):
+        elif not is_valid_inferred_type(init_type) and not self.no_partial_types:
             # We cannot use the type of the initialization expression for full type
             # inference (it's not specific enough), but we might be able to give
             # partial type which will be made more specific later. A partial type
@@ -1897,7 +1965,7 @@ def check_member_assignment(self, instance_type: Type, attribute_type: Type,
                                 rvalue: Expression, context: Context) -> Tuple[Type, bool]:
         """Type member assigment.
 
-        This is defers to check_simple_assignment, unless the member expression
+        This defers to check_simple_assignment, unless the member expression
         is a descriptor, in which case this checks descriptor semantics as well.
 
         Return the inferred rvalue_type and whether to infer anything about the attribute type
@@ -2697,7 +2765,19 @@ def iterable_item_type(self, instance: Instance) -> Type:
         iterable = map_instance_to_supertype(
             instance,
             self.lookup_typeinfo('typing.Iterable'))
-        return iterable.args[0]
+        item_type = iterable.args[0]
+        if not isinstance(item_type, AnyType):
+            # This relies on 'map_instance_to_supertype' returning 'Iterable[Any]'
+            # in case there is no explicit base class.
+            return item_type
+        # Try also structural typing.
+        iter_type = find_member('__iter__', instance, instance)
+        if (iter_type and isinstance(iter_type, CallableType) and
+                isinstance(iter_type.ret_type, Instance)):
+            iterator = map_instance_to_supertype(iter_type.ret_type,
+                                                 self.lookup_typeinfo('typing.Iterator'))
+            item_type = iterator.args[0]
+        return item_type
 
     def function_type(self, func: FuncBase) -> FunctionLike:
         return function_type(func, self.named_type('builtins.function'))

mypy/maptype.py

@@ -10,7 +10,8 @@ def map_instance_to_supertype(instance: Instance,
     """Produce a supertype of `instance` that is an Instance
     of `superclass`, mapping type arguments up the chain of bases.
 
-    `superclass` is required to be a superclass of `instance.type`.
+    If `superclass` is not a nominal superclass of `instance.type`,
+    then all type arguments are mapped to 'Any'.
     """
     if instance.type == superclass:
         # Fast path: `instance` already belongs to `superclass`.