working implementation

python · JukkaL · Mar 18, 2017 · Jan 17, 2017 · Jan 18, 2017 · Jan 18, 2017
commit c64e7fc63096b317c97ca5911c92ebabaebe48d8
diff --git a/mypy/checker.py b/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
@@ -357,7 +358,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()
@@ -388,7 +390,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()
@@ -402,6 +405,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.

diff --git a/mypy/semanal.py b/mypy/semanal.py
@@ -1294,7 +1294,6 @@ def normalize_type_alias(self, node: SymbolTableNode,
         normalized = False
         if node.fullname in type_aliases:<
8000
/span>
             # Node refers to an aliased type such as typing.List; normalize.
-            old_node = node
             node = self.lookup_qualified(type_aliases[node.fullname], ctx)
             normalized = True
         if node.fullname in collections_type_aliases:

diff --git a/test-data/unit/check-async-await.test b/test-data/unit/check-async-await.test
@@ -410,25 +410,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
 -- ------------------------------------------

diff --git a/test-data/unit/fixtures/dict.pyi b/test-data/unit/fixtures/dict.pyi
@@ -38,3 +38,5 @@ class tuple: pass
 class function: pass
 class float: pass
 class bool: pass
+
+class BaseException: pass