Fine-grained: Support async def/for/with and await (#4969)

JukkaL · web-flow · commit cdfc272fafab · 2018-04-26T12:21:39.000+01:00
Some of the changes motivated cleaning up some duplicate error messages. Also noticed missing __next__/next dependency in ordinary for statements and fixed it as well. Work towards #4951.
diff --git a/mypy/checker.py b/mypy/checker.py
@@ -2517,57 +2517,47 @@ def get_types_from_except_handler(self, typ: Type, n: Expression) -> List[Type]:
     def visit_for_stmt(self, s: ForStmt) -> None:
         """Type check a for statement."""
         if s.is_async:
-            item_type = self.analyze_async_iterable_item_type(s.expr)
+            iterator_type, item_type = self.analyze_async_iterable_item_type(s.expr)
         else:
-            item_type = self.analyze_iterable_item_type(s.expr)
+            iterator_type, item_type = self.analyze_iterable_item_type(s.expr)
         s.inferred_item_type = item_type
+        s.inferred_iterator_type = iterator_type
         self.analyze_index_variables(s.index, item_type, s.index_type is None, s)
         self.accept_loop(s.body, s.else_body)
 
-    def analyze_async_iterable_item_type(self, expr: Expression) -> Type:
-        """Analyse async iterable expression and return iterator item type."""
+    def analyze_async_iterable_item_type(self, expr: Expression) -> Tuple[Type, Type]:
+        """Analyse async iterable expression and return iterator and iterator item types."""
         echk = self.expr_checker
         iterable = echk.accept(expr)
-
-        self.check_subtype(iterable,
-                           self.named_generic_type('typing.AsyncIterable',
-                                                   [AnyType(TypeOfAny.special_form)]),
-                           expr, messages.ASYNC_ITERABLE_EXPECTED)
-
         method = echk.analyze_external_member_access('__aiter__', iterable, expr)
         iterator = echk.check_call(method, [], [], expr)[0]
         method = echk.analyze_external_member_access('__anext__', iterator, expr)
         awaitable = echk.check_call(method, [], [], expr)[0]
-        return echk.check_awaitable_expr(awaitable, expr,
-                                         messages.INCOMPATIBLE_TYPES_IN_ASYNC_FOR)
+        item_type = echk.check_awaitable_expr(awaitable, expr,
+                                              messages.INCOMPATIBLE_TYPES_IN_ASYNC_FOR)
+        return iterator, item_type
 
-    def analyze_iterable_item_type(self, expr: Expression) -> Type:
-        """Analyse iterable expression and return iterator item type."""
+    def analyze_iterable_item_type(self, expr: Expression) -> Tuple[Type, Type]:
+        """Analyse iterable expression and return iterator and iterator item types."""
         echk = self.expr_checker
         iterable = echk.accept(expr)
+        method = echk.analyze_external_member_access('__iter__', iterable, expr)
+        iterator = echk.check_call(method, [], [], expr)[0]
 
         if isinstance(iterable, TupleType):
             joined = UninhabitedType()  # type: Type
             for item in iterable.items:
                 joined = join_types(joined, item)
-            return joined
+            return iterator, joined
         else:
             # Non-tuple iterable.
-            self.check_subtype(iterable,
-                               self.named_generic_type('typing.Iterable',
-                                                       [AnyType(TypeOfAny.special_form)]),
-                               expr, messages.ITERABLE_EXPECTED)
-
-            method = echk.analyze_external_member_access('__iter__', iterable,
-                                                         expr)
-            iterator = echk.check_call(method, [], [], expr)[0]
             if self.options.python_version[0] >= 3:
                 nextmethod = '__next__'
             else:
                 nextmethod = 'next'
             method = echk.analyze_external_member_access(nextmethod, iterator,
                                                          expr)
-            return echk.check_call(method, [], [], expr)[0]
+            return iterator, echk.check_call(method, [], [], expr)[0]
 
     def analyze_index_variables(self, index: Expression, item_type: Type,
                                 infer_lvalue_type: bool, context: Context) -> None:
diff --git a/mypy/checkexpr.py b/mypy/checkexpr.py
@@ -1422,7 +1422,7 @@ def visit_comparison_expr(self, e: ComparisonExpr) -> Type:
                 elif (local_errors.is_errors() and
                     # is_valid_var_arg is True for any Iterable
                         self.is_valid_var_arg(right_type)):
-                    itertype = self.chk.analyze_iterable_item_type(right)
+                    _, itertype = self.chk.analyze_iterable_item_type(right)
                     method_type = CallableType(
                         [left_type],
                         [nodes.ARG_POS],
@@ -2290,9 +2290,9 @@ def check_for_comp(self, e: Union[GeneratorExpr, DictionaryComprehension]) -> No
         for index, sequence, conditions, is_async in zip(e.indices, e.sequences,
                                                          e.condlists, e.is_async):
             if is_async:
-                sequence_type = self.chk.analyze_async_iterable_item_type(sequence)
+                _, sequence_type = self.chk.analyze_async_iterable_item_type(sequence)
             else:
-                sequence_type = self.chk.analyze_iterable_item_type(sequence)
+                _, sequence_type = self.chk.analyze_iterable_item_type(sequence)
             self.chk.analyze_index_variables(index, sequence_type, True, e)
             for condition in conditions:
                 self.accept(condition)
diff --git a/mypy/messages.py b/mypy/messages.py
@@ -63,8 +63,6 @@
 MUST_HAVE_NONE_RETURN_TYPE = 'The return type of "{}" must be None'
 INVALID_TUPLE_INDEX_TYPE = 'Invalid tuple index type'
 TUPLE_INDEX_OUT_OF_RANGE = 'Tuple index out of range'
-ITERABLE_EXPECTED = 'Iterable expected'
-ASYNC_ITERABLE_EXPECTED = 'AsyncIterable expected'
 INVALID_SLICE_INDEX = 'Slice index must be an integer or None'
 CANNOT_INFER_LAMBDA_TYPE = 'Cannot infer type of lambda'
 CANNOT_INFER_ITEM_TYPE = 'Cannot infer iterable item type'
@@ -450,28 +448,36 @@ def has_no_attr(self, original_type: Type, typ: Type, member: str, context: Cont
                 self.fail('{} not callable'.format(self.format(original_type)), context)
         else:
             # The non-special case: a missing ordinary attribute.
+            extra = ''
+            if member == '__iter__':
+                extra = ' (not iterable)'
+            elif member == '__aiter__':
+                extra = ' (not async iterable)'
             if not self.disable_type_names:
                 failed = False
                 if isinstance(original_type, Instance) and original_type.type.names:
                     alternatives = set(original_type.type.names.keys())
                     matches = [m for m in COMMON_MISTAKES.get(member, []) if m in alternatives]
                     matches.extend(best_matches(member, alternatives)[:3])
+                    if member == '__aiter__' and matches == ['__iter__']:
+                        matches = []  # Avoid misleading suggestion
                     if matches:
-                        self.fail('{} has no attribute "{}"; maybe {}?'.format(
-                            self.format(original_type), member, pretty_or(matches)), context)
+                        self.fail('{} has no attribute "{}"; maybe {}?{}'.format(
+                            self.format(original_type), member, pretty_or(matches), extra),
+                            context)
                         failed = True
                 if not failed:
-                    self.fail('{} has no attribute "{}"'.format(self.format(original_type),
-                                                                member), context)
+                    self.fail('{} has no attribute "{}"{}'.format(self.format(original_type),
+                                                                  member, extra), context)
             elif isinstance(original_type, UnionType):
                 # The checker passes "object" in lieu of "None" for attribute
                 # checks, so we manually convert it back.
                 typ_format = self.format(typ)
                 if typ_format == '"object"' and \
                         any(type(item) == NoneTyp for item in original_type.items):
                     typ_format = '"None"'
-                self.fail('Item {} of {} has no attribute "{}"'.format(
-                    typ_format, self.format(original_type), member), context)
+                self.fail('Item {} of {} has no attribute "{}"{}'.format(
+                    typ_format, self.format(original_type), member, extra), context)
         return AnyType(TypeOfAny.from_error)
 
     def unsupported_operand_types(self, op: str, left_type: Any,
diff --git a/mypy/nodes.py b/mypy/nodes.py
@@ -952,6 +952,8 @@ class ForStmt(Statement):
     index_type = None  # type: Optional[mypy.types.Type]
     # Inferred iterable item type
     inferred_item_type = None  # type: Optional[mypy.types.Type]
+    # Inferred iterator type
+    inferred_iterator_type = None  # type: Optional[mypy.types.Type]
     # Expression to iterate
     expr = None  # type: Expression
     body = None  # type: Block
diff --git a/mypy/server/aststrip.py b/mypy/server/aststrip.py
@@ -44,7 +44,7 @@
     Node, FuncDef, NameExpr, MemberExpr, RefExpr, MypyFile, FuncItem, ClassDef, AssignmentStmt,
     ImportFrom, Import, TypeInfo, SymbolTable, Var, CallExpr, Decorator, OverloadedFuncDef,
     SuperExpr, UNBOUND_IMPORTED, GDEF, MDEF, IndexExpr, SymbolTableNode, ImportAll, TupleExpr,
-    ListExpr
+    ListExpr, ForStmt
 )
 from mypy.semanal_shared import create_indirect_imported_name
 from mypy.traverser import TraverserVisitor
@@ -217,6 +217,11 @@ def visit_import(self, node: Import) -> None:
                     symnode.kind = UNBOUND_IMPORTED
                     symnode.node = None
 
+    def visit_for_stmt(self, node: ForStmt) -> None:
+        node.index_type = None
+        node.inferred_item_type = None
+        super().visit_for_stmt(node)
+
     def visit_import_all(self, node: ImportAll) -> None:
         # If the node is unreachable, we don't want to reset entries from a reachable import.
         if node.is_unreachable:
diff --git a/mypy/server/deps.py b/mypy/server/deps.py
@@ -92,7 +92,7 @@ class 'mod.Cls'. This can also refer to an attribute inherited from a
     ComparisonExpr, GeneratorExpr, DictionaryComprehension, StarExpr, PrintStmt, ForStmt, WithStmt,
     TupleExpr, ListExpr, OperatorAssignmentStmt, DelStmt, YieldFromExpr, Decorator, Block,
     TypeInfo, FuncBase, OverloadedFuncDef, RefExpr, SuperExpr, Var, NamedTupleExpr, TypedDictExpr,
-    LDEF, MDEF, GDEF, FuncItem, TypeAliasExpr, NewTypeExpr, ImportAll, EnumCallExpr,
+    LDEF, MDEF, GDEF, FuncItem, TypeAliasExpr, NewTypeExpr, ImportAll, EnumCallExpr, AwaitExpr,
     op_methods, reverse_op_methods, ops_with_inplace_method, unary_op_methods
 )
 from mypy.traverser import TraverserVisitor
@@ -166,7 +166,6 @@ def __init__(self,
         self.is_package_init_file = False
 
     # TODO (incomplete):
-    #   await
     #   protocols
 
     def visit_mypy_file(self, o: MypyFile) -> None:
@@ -415,10 +414,22 @@ def visit_operator_assignment_stmt(self, o: OperatorAssignmentStmt) -> None:
 
     def visit_for_stmt(self, o: ForStmt) -> None:
         super().visit_for_stmt(o)
-        # __getitem__ is only used if __iter__ is missing but for simplicity we
-        # just always depend on both.
-        self.add_attribute_dependency_for_expr(o.expr, '__iter__')
-        self.add_attribute_dependency_for_expr(o.expr, '__getitem__')
+        if not o.is_async:
+            # __getitem__ is only used if __iter__ is missing but for simplicity we
+            # just always depend on both.
+            self.add_attribute_dependency_for_expr(o.expr, '__iter__')
+            self.add_attribute_dependency_for_expr(o.expr, '__getitem__')
+            if o.inferred_iterator_type:
+                if self.python2:
+                    method = 'next'
+                else:
+                    method = '__next__'
+                self.add_attribute_dependency(o.inferred_iterator_type, method)
+        else:
+            self.add_attribute_dependency_for_expr(o.expr, '__aiter__')
+            if o.inferred_iterator_type:
+                self.add_attribute_dependency(o.inferred_iterator_type, '__anext__')
+
         self.process_lvalue(o.index)
         if isinstance(o.index, TupleExpr):
             # Process multiple assignment to index variables.
@@ -433,8 +444,12 @@ def visit_for_stmt(self, o: ForStmt) -> None:
     def visit_with_stmt(self, o: WithStmt) -> None:
         super().visit_with_stmt(o)
         for e in o.expr:
-            self.add_attribute_dependency_for_expr(e, '__enter__')
-            self.add_attribute_dependency_for_expr(e, '__exit__')
+            if not o.is_async:
+                self.add_attribute_dependency_for_expr(e, '__enter__')
+                self.add_attribute_dependency_for_expr(e, '__exit__')
+            else:
+                self.add_attribute_dependency_for_expr(e, '__aenter__')
+                self.add_attribute_dependency_for_expr(e, '__aexit__')
         if o.target_type:
             self.add_type_dependencies(o.target_type)
 
@@ -584,6 +599,10 @@ def visit_yield_from_expr(self, e: YieldFromExpr) -> None:
         super().visit_yield_from_expr(e)
         self.add_iter_dependency(e.expr)
 
+    def visit_await_expr(self, e: AwaitExpr) -> None:
+        super().visit_await_expr(e)
+        self.add_attribute_dependency_for_expr(e.expr, '__await__')
+
     # Helpers
 
     def add_type_alias_deps(self, target: str) -> None:
diff --git a/test-data/unit/check-async-await.test b/test-data/unit/check-async-await.test
@@ -163,8 +163,7 @@ async def f() -> None:
 [builtins fixtures/async_await.pyi]
 [typing fixtures/typing-full.pyi]
 [out]
-main:4: error: AsyncIterable expected
-main:4: error: "List[int]" has no attribute "__aiter__"
+main:4: error: "List[int]" has no attribute "__aiter__" (not async iterable)
 
 [case testAsyncForTypeComments]
 
@@ -247,14 +246,10 @@ async def wrong_iterable(obj: Iterable[int]):
     {i: i for i in asyncify(obj)}
 
 [out]
-main:18: error: AsyncIterable expected
-main:18: error: "Iterable[int]" has no attribute "__aiter__"; maybe "__iter__"?
-main:19: error: Iterable expected
-main:19: error: "asyncify[int]" has no attribute "__iter__"; maybe "__aiter__"?
-main:20: error: AsyncIterable expected
-main:20: error: "Iterable[int]" has no attribute "__aiter__"; maybe "__iter__"?
-main:21: error: Iterable expected
-main:21: error: "asyncify[int]" has no attribute "__iter__"; maybe "__aiter__"?
+main:18: error: "Iterable[int]" has no attribute "__aiter__" (not async iterable)
+main:19: error: "asyncify[int]" has no attribute "__iter__"; maybe "__aiter__"? (not iterable)
+main:20: error: "Iterable[int]" has no attribute "__aiter__" (not async iterable)
+main:21: error: "asyncify[int]" has no attribute "__iter__"; maybe "__aiter__"? (not iterable)
 [builtins fixtures/async_await.pyi]
 [typing fixtures/typing-full.pyi]
 
@@ -538,7 +533,7 @@ async def h() -> None:
 from typing import AsyncGenerator
 
 async def gen() -> AsyncGenerator[int, None]:
-    for i in (1, 2, 3):
+    for i in [1, 2, 3]:
         yield i
 
 def h() -> None:
@@ -549,8 +544,7 @@ def h() -> None:
 [typing fixtures/typing-full.pyi]
 
 [out]
-main:9: error: Iterable expected
-main:9: error: "AsyncGenerator[int, None]" has no attribute "__iter__"; maybe "__aiter__"?
+main:9: error: "AsyncGenerator[int, None]" has no attribute "__iter__"; maybe "__aiter__"? (not iterable)
 
 [case testAsyncGeneratorNoYieldFrom]
 # flags: --fast-parser --python-version 3.6
diff --git a/test-data/unit/check-incomplete-fixture.test b/test-data/unit/check-incomplete-fixture.test
@@ -76,8 +76,7 @@ for y in x:
 -- avoid things getting worse.
 main:2: error: "tuple" expects no type arguments, but 1 given
 main:3: error: Value of type "tuple" is not indexable
-main:4: error: Iterable expected
-main:4: error: "tuple" has no attribute "__iter__"
+main:4: error: "tuple" has no attribute "__iter__" (not iterable)
 
 [case testClassmethodMissingFromStubs]
 class A:
diff --git a/test-data/unit/check-inference.test b/test-data/unit/check-inference.test
@@ -1563,7 +1563,7 @@ class A:
             pass
 [builtins fixtures/for.pyi]
 [out]
-main:5: error: "None" has no attribute "__iter__"
+main:5: error: "None" has no attribute "__iter__" (not iterable)
 
 [case testPartialTypeErrorSpecialCase2]
 # This used to crash.
@@ -1584,7 +1584,7 @@ class A:
             pass
 [builtins fixtures/for.pyi]
 [out]
-main:4: error: "None" has no attribute "__iter__"
+main:4: error: "None" has no attribute "__iter__" (not iterable)
 
 
 -- Multipass
diff --git a/test-data/unit/check-unions.test b/test-data/unit/check-unions.test
@@ -869,8 +869,7 @@ a: Any
 d: Dict[str, Tuple[List[Tuple[str, str]], str]]
 x, _ = d.get(a, (None, None))
 
-for y in x: pass # E: Iterable expected \
-                 # E: Item "None" of "Optional[List[Tuple[str, str]]]" has no attribute "__iter__"
+for y in x: pass # E: Item "None" of "Optional[List[Tuple[str, str]]]" has no attribute "__iter__" (not iterable)
 if x:
     for s, t in x:
         reveal_type(s) # E: Revealed type is 'builtins.str'
@@ -886,8 +885,7 @@ x = None
 d: Dict[str, Tuple[List[Tuple[str, str]], str]]
 x, _ = d.get(a, (None, None))
 
-for y in x: pass # E: Iterable expected \
-                 # E: Item "None" of "Optional[List[Tuple[str, str]]]" has no attribute "__iter__"
+for y in x: pass # E: Item "None" of "Optional[List[Tuple[str, str]]]" has no attribute "__iter__" (not iterable)
 if x:
     for s, t in x:
         reveal_type(s) # E: Revealed type is 'builtins.str'
diff --git a/test-data/unit/deps.test b/test-data/unit/deps.test
@@ -1025,6 +1025,49 @@ class B(Base):
 <mod.Base> -> <m.f>, m.f
 <mod> -> m
 
+[case testCustomIterator]
+class A:
+    def __iter__(self) -> B: pass
+class B:
+    def __iter__(self) -> B: pass
+    def __next__(self) -> C: pass
+class C:
+    pass
+def f() -> None:
+    for x in A(): pass
+[out]
+<m.A.__getitem__> -> m.f
+<m.A.__init__> -> m.f
+<m.A.__iter__> -> m.f
+<m.A.__new__> -> m.f
+<m.A> -> m.A, m.f
+<m.B.__next__> -> m.f
+<m.B> -> <m.A.__iter__>, <m.B.__iter__>, m.A.__iter__, m.B, m.B.__iter__
+<m.C> -> <m.B.__next__>, m.B.__next__, m.C
+
+[case testCustomIterator_python2]
+class A:
+    def __iter__(self): # type: () -> B
+        pass
+class B:
+    def __iter__(self): # type: () -> B
+        pass
+    def next(self): # type: () -> C
+        pass
+class C:
+    pass
+def f(): # type: () -> None
+    for x in A(): pass
+[out]
+<m.A.__getitem__> -> m.f
+<m.A.__init__> -> m.f
+<m.A.__iter__> -> m.f
+<m.A.__new__> -> m.f
+<m.A> -> m.A, m.f
+<m.B.next> -> m.f
+<m.B> -> <m.A.__iter__>, <m.B.__iter__>, m.A.__iter__, m.B, m.B.__iter__
+<m.C> -> <m.B.next>, m.B.next, m.C
+
 [case testDepsLiskovClass]
 from mod import A, C
 class D(C):
diff --git a/test-data/unit/fine-grained.test b/test-data/unit/fine-grained.test
