working implementation

Author: JelleZijlstra

mypy/checker.py

@@ -3,6 +3,7 @@
 import itertools
 import fnmatch
 from contextlib import contextmanager
+import sys
 
 from typing import (
     Dict, Set, List, cast, Tuple, TypeVar, Union, Optional, NamedTuple, Iterator
@@ -380,7 +381,8 @@ def get_generator_yield_type(self, return_type: Type, is_coroutine: bool) -> Typ
         """Given the declared return type of a generator (t), return the type it yields (ty)."""
         if isinstance(return_type, AnyType):
             return AnyType()
-        elif not self.is_generator_return_type(return_type, is_coroutine):
+        elif (not self.is_generator_return_type(return_type, is_coroutine)
+                and not self.is_async_generator_return_type(return_type)):
             # If the function doesn't have a proper Generator (or
             # Awaitable) return type, anything is permissible.
             return AnyType()
@@ -411,7 +413,8 @@ def get_generator_receive_type(self, return_type: Type, is_coroutine: bool) -> T
         """Given a declared generator return type (t), return the type its yield receives (tc)."""
         if isinstance(return_type, AnyType):
             return AnyType()
-        elif not self.is_generator_return_type(return_type, is_coroutine):
+        elif (not self.is_generator_return_type(return_type, is_coroutine)
+                and not self.is_async_generator_return_type(return_type)):
             # If the function doesn't have a proper Generator (or
             # Awaitable) return type, anything is permissible.
             return AnyType()
@@ -425,6 +428,8 @@ def get_generator_receive_type(self, return_type: Type, is_coroutine: bool) -> T
               and len(return_type.args) >= 3):
             # Generator: tc is args[1].
             return return_type.args[1]
+        elif return_type.type.fullname() == 'typing.AsyncGenerator' and len(return_type.args) >= 2:
+            return return_type.args[1]
         else:
             # `return_type` is a supertype of Generator, so callers won't be able to send it
             # values.  IOW, tc is None.

mypy/semanal.py

@@ -1068,9 +1068,7 @@ def normalize_type_alias(self, node: SymbolTableNode,
                              ctx: Context) -> SymbolTableNode:
         if node.fullname in type_aliases:
             # Node refers to an aliased type such as typing.List; normalize.
-            old_node = node
             node = self.lookup_qualified(type_aliases[node.fullname], ctx)
-            assert node is not None, (type_aliases, old_node, old_node.fullname, ctx)
         if node.fullname == 'typing.DefaultDict':
             self.add_module_symbol('collections', '__mypy_collections__', False, ctx)
             node = self.lookup_qualified('__mypy_collections__.defaultdict', ctx)

test-data/unit/check-async-await.test

@@ -303,25 +303,125 @@ def f() -> Generator[int, str, int]:
 [builtins fixtures/async_await.pyi]
 [out]
 
+-- Async generators (PEP 525), some test cases adapted from the PEP text
+-- ---------------------------------------------------------------------
+
 [case testAsyncGenerator]
 # flags: --fast-parser --python-version 3.6
 from mypy_extensions import AsyncGenerator
 
 async def f() -> int:
     return 42
 
-async def g() -> AsyncGenerator[int, str]:
+async def g() -> AsyncGenerator[int, None]:
     value = await f()
-    reveal_type(value)  # E: Revealed type is 'builtins.int'
-    x = yield value
-    reveal_type(x)  # E: Revealed type is 'builtins.str'
+    reveal_type(value)  # E: Revealed type is 'builtins.int*'
+    yield value
+    # return without a value is fine
+    return
+reveal_type(g)  # E: Revealed type is 'def () -> typing.AsyncGenerator[builtins.int, void]'
+reveal_type(g())  # E: Revealed type is 'typing.AsyncGenerator[builtins.int, void]'
 
 async def h() -> None:
     async for item in g():
-        reveal_type(item)  # E: Revealed type is 'builtins.int'
+        reveal_type(item)  # E: Revealed type is 'builtins.int*'
+
+[builtins fixtures/dict.pyi]
+
+[case testAsyncGeneratorManualIter]
+# flags: --fast-parser --python-version 3.6
+from mypy_extensions import AsyncGenerator
+
+async def genfunc() -> AsyncGenerator[int, None]:
+    yield 1
+    yield 2
+
+async def user() -> None:
+    gen = genfunc()
+
+    reveal_type(gen.__aiter__())  # E: Revealed type is 'typing.AsyncGenerator[builtins.int*, void]'
+
+    reveal_type(await gen.__anext__())  # E: Revealed type is 'builtins.int*'
+
+[builtins fixtures/dict.pyi]
+
+[case testAsyncGeneratorAsend]
+# flags: --fast-parser --python-version 3.6
+from mypy_extensions import AsyncGenerator
+
+async def f() -> None:
+    pass
+
+async def gen() -> AsyncGenerator[int, str]:
+    await f()
+    v = yield 42
+    reveal_type(v)  # E: Revealed type is 'builtins.str'
+    await f()
+
+async def h() -> None:
+    g = gen()
+    await g.asend(())  # E: Argument 1 to "asend" of "AsyncGenerator" has incompatible type "Tuple[]"; expected "str"
+    reveal_type(await g.asend('hello'))  # E: Revealed type is 'builtins.int*'
 
 [builtins fixtures/dict.pyi]
 
+[case testAsyncGeneratorAthrow]
+# flags: --fast-parser --python-version 3.6
+from mypy_extensions import AsyncGenerator
+
+async def gen() -> AsyncGenerator[str, int]:
+    try:
+        yield 'hello'
+    except BaseException:
+        yield 'world'
+
+async def h() -> None:
+    g = gen()
+    v = await g.asend(1)
+    reveal_type(v)  # E: Revealed type is 'builtins.str*'
+    await g.athrow(BaseException)
+
+[builtins fixtures/dict.pyi]
+
+[case testAsyncGeneratorNoSyncIteration]
+# flags: --fast-parser --python-version 3.6
+from mypy_extensions import AsyncGenerator
+
+async def gen() -> AsyncGenerator[int, None]:
+    for i in (1, 2, 3):
+        yield i
+
+def h() -> None:
+    for i in gen():
+        pass
+
+[builtins fixtures/dict.pyi]
+
+[out]
+main:9: error: Iterable expected
+main:9: error: AsyncGenerator[int, None] has no attribute "__iter__"; maybe "__aiter__"?
+
+[case testAsyncGeneratorNoYieldFrom]
+# flags: --fast-parser --python-version 3.6
+from mypy_extensions import AsyncGenerator
+
+async def f() -> AsyncGenerator[int, None]:
+    pass
+
+async def gen() -> AsyncGenerator[int, None]:
+    yield from f()  # E: 'yield from' in async function
+
+[builtins fixtures/dict.pyi]
+
+[case testAsyncGeneratorNoReturnWithValue]
+# flags: --fast-parser --python-version 3.6
+from mypy_extensions import AsyncGenerator
+
+async def gen() -> AsyncGenerator[int, None]:
+    yield 1
+    return 42
+
+[builtins fixtures/dict.pyi]
 
 -- The full matrix of coroutine compatibility
 -- ------------------------------------------

test-data/unit/fixtures/dict.pyi

@@ -33,3 +33,5 @@ class list(Iterable[T], Generic[T]): # needed by some test cases
 class tuple: pass
 class function: pass
 class float: pass
+
+class BaseException: pass