Adapt some old work

python · JukkaL · Oct 11, 2017 · Jul 24, 2017 · Oct 4, 2017 · Oct 5, 2017
commit d550e4ae2d244534cec8b49b77b7a7f073dbdbc7
diff --git a/mypy/binder.py b/mypy/binder.py
@@ -1,5 +1,6 @@
-from typing import Dict, List, Set, Iterator, Union, Optional, cast
+from typing import Dict, List, Set, Iterator, Union, Optional, Tuple, DefaultDict, cast
 from contextlib import contextmanager
+from collections import defaultdict
 
 from mypy.types import Type, AnyType, PartialType, UnionType, TypeOfAny
 from mypy.subtypes import is_subtype
@@ -57,6 +58,7 @@ class A:
     reveal_type(lst[0].a) # str
     ```
     """
+    type_assignments = None  # type: Optional[DefaultDict[Expression, List[Tuple[Type, Type]]]]
 
     def __init__(self) -> None:
         # The stack of frames currently used.  These map
@@ -210,10 +212,20 @@ def pop_frame(self, can_skip: bool, fall_through: int) -> Frame:
 
         return result
 
+    @contextmanager
+    def accumulate_type_assignments(self) -> Iterator[DefaultDict[Expression,
+                                                                  List[Tuple[Type, Type]]]]:
+        self.type_assignments = defaultdict(list)
+        yield self.type_assignments
+        self.type_assignments = None
+
     def assign_type(self, expr: Expression,
                     type: Type,
                     restrict_any: bool = False) -> None:
+        if self.type_assignments is not None:
+            self.type_assignments[expr].append((type, declared_type))
+            return
         if not isinstance(expr, BindableTypes):
             return None
         if not literal(expr):

diff --git a/mypy/checker.py b/mypy/checker.py
@@ -53,7 +53,7 @@
 from mypy.erasetype import erase_typevars
 from mypy.expandtype import expand_type, expand_type_by_instance
 from mypy.visitor import NodeVisitor
-from mypy.join import join_types
+from mypy.join import join_types, join_type_list
 from mypy.treetransform import TransformVisitor
 from mypy.binder import ConditionalTypeBinder, get_declaration
 from mypy.meet import is_overlapping_types
@@ -1604,12 +1604,14 @@ def check_rvalue_count_in_assignment(self, lvalues: List[Lvalue], rvalue_count:
     def check_multi_assignment(self, lvalues: List[Lvalue],
                                rvalue: Expression,
                                context: Context,
-                               infer_lvalue_type: bool = True) -> None:
+                               infer_lvalue_type: bool = True,
+                               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,11 +1630,37 @@ def check_multi_assignment(self, lvalues: List[Lvalue],
             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, context, infer_lvalue_type)
+            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:
+        transposed = tuple([] for _ in lvalues)  # type: Tuple[List[Type], ...]
+        with self.binder.accumulate_type_assignments() as assignments:
+            for item in rvalue_type.items:
+                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):
+                    t.append(self.type_map.get(lv, AnyType(TypeOfAny.special_form)))
+        union_types = tuple(join_type_list(col) for col in transposed)
+        for expr, items in assignments.items():
+            types, declared_types = zip(*items)
+            self.binder.assign_type(expr,
+                                    join_type_list(types),
+                                    join_type_list(declared_types),
+                                    False)
+        for union, lv in zip(union_types, lvalues):
+            _1, _2, inferred = self.check_lvalue(lv)
+            if inferred:
+                self.set_inferred_type(inferred, lv, union)
+            else:
+                self.store_type(lv, union)
+
     def check_multi_assignment_from_tuple(self, lvalues: List[Lvalue], rvalue: Expression,
                                           rvalue_type: TupleType, context: Context,
                                           undefined_rvalue: bool,
@@ -1820,7 +1848,6 @@ def infer_variable_type(self, name: Var, lvalue: Lvalue,
 
             # Make the type more general (strip away function names etc.).
             init_type = strip_type(init_type)
-
             self.set_inferred_type(name, lvalue, init_type)
 
     def infer_partial_type(self, name: Var, lvalue: Lvalue, init_type: Type) -> bool:
@@ -2698,7 +2725,17 @@ 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):
+            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'))

diff --git a/test-data/unit/fixtures/list.pyi b/test-data/unit/fixtures/list.pyi
@@ -19,6 +19,7 @@ class list(Generic[T]):
     def __add__(self, x: list[T]) -> list[T]: pass
     def __mul__(self, x: int) -> list[T]: pass
     def __getitem__(self, x: int) -> T: pass
+    def __setitem__(self, x: int, v: T) -> None: pass
     def append(self, x: T) -> None: pass
     def extend(self, x: Iterable[T]) -> None: pass