Repository: quora/pyanalyze
Pull request #594 (2023-02-03 14:31:16)

Better support for async generators


Patch:
diff --git a/docs/changelog.md b/docs/changelog.md
index 4ed6676..026f4cc 100644
--- a/docs/changelog.md
+++ b/docs/changelog.md
@@ -1,13 +1,19 @@
+# Unreleased
+
+- Better typechecking support for async generators (#594)
+
 # Changelog
 
 ## Version 0.9.0 (January 16, 2023)
 
 Release highlights:
+
 - Support for PEP 702 (`@typing.deprecated`) (#578)
 - Add experimental `@has_extra_keys` decorator for `TypedDict` types
 - Support more Python 3.11 features (`except*` and `get_overloads`)
 
 Full changelog:
+
 - Support `typing_extensions.get_overloads` and `typing.get_overloads` (#589)
 - Support `in` on objects with only `__iter__` (#588)
 - Do not call `.mro()` method on non-types (#587)
@@ -26,11 +32,13 @@ Full changelog:
 ## Version 0.8.0 (November 5, 2022)
 
 Release highlights:
+
 - Support for Python 3.11
 - Drop support for Python 3.6
-- Support for PEP 692 (``Unpack`` on ``**kwargs``)
+- Support for PEP 692 (`Unpack` on `**kwargs`)
 
 Full changelog:
+
 - Infer `async def` functions as returning `Coroutine`, not
   `Awaitable` (#557, #559)
 - Drop support for Python 3.6 (#554)
@@ -63,11 +71,13 @@ Full changelog:
 ## Version 0.7.0 (April 13, 2022)
 
 Release highlights:
+
 - Support for PEP 673 (`Self`)
 - Support for PEP 675 (`LiteralString`)
 - Support for `assert_type` and other additions to `typing` in Python 3.11
 
 Full changelog:
+
 - Remove `SequenceIncompleteValue` (#519)
 - Add implementation function for `dict.pop` (#517)
 - Remove `WeakExtension` (#517)
@@ -115,7 +125,7 @@ Full changelog:
 - Do not treat attributes that raise an exception on access
   as nonexistent (#481)
 - Improve detection of unhashable dict keys and set members (#469)
-- The `in` and `not in` operators always return 
+- The `in` and `not in` operators always return
   booleans (#480)
 - Allow `NotImplemented` to be returned from special
   methods that support it (#479)
@@ -166,13 +176,14 @@ Full changelog:
 - Handle `NoReturn` in `async def` functions (#427)
 - Support PEP 673 (`typing_extensions.Self`) (#423)
 - Updates for compatibility with recent changes in typeshed (#421):
-  - Fix override compatibility check for unknown callables 
+  - Fix override compatibility check for unknown callables
   - Fix usage of removed type `_typeshed.SupportsLessThan`
 - Remove old configuration abstraction (#414)
 
 ## Version 0.6.0 (January 12, 2022)
 
 Release highlights:
+
 - Support for configuration through `pyproject.toml`. The old
   configuration mechanism will be removed in the next release.
 - Support for experimental new type evaluation mechanism, providing
@@ -180,6 +191,7 @@ Release highlights:
 - Support for suggesting annotations for unannotated code.
 
 Full changelog:
+
 - Support generic type evaluators (#409)
 - Implement return annotation behavior for type evaluation
   functions (#408)
diff --git a/pyanalyze/functions.py b/pyanalyze/functions.py
index 9edf6b1..b5a74c9 100644
--- a/pyanalyze/functions.py
+++ b/pyanalyze/functions.py
@@ -32,6 +32,7 @@ from .value import (
     GenericValue,
     get_tv_map,
     KnownValue,
+    is_async_iterable,
     is_iterable,
     make_coro_type,
     SubclassValue,
@@ -53,6 +54,9 @@ GeneratorValue = GenericValue(
     collections.abc.Generator,
     [TypeVarValue(YieldT), TypeVarValue(SendT), TypeVarValue(ReturnT)],
 )
+AsyncGeneratorValue = GenericValue(
+    collections.abc.AsyncGenerator, [TypeVarValue(YieldT), TypeVarValue(SendT)]
+)
 
 
 class AsyncFunctionKind(enum.Enum):
@@ -92,23 +96,40 @@ class FunctionInfo:
     def get_generator_yield_type(self, ctx: CanAssignContext) -> Value:
         if self.return_annotation is None:
             return AnyValue(AnySource.unannotated)
-        iterable_val = is_iterable(self.return_annotation, ctx)
-        if isinstance(iterable_val, CanAssignError):
-            return AnyValue(AnySource.error)
-        return iterable_val
+        if isinstance(self.node, ast.AsyncFunctionDef):
+            iterable_val = is_async_iterable(self.return_annotation, ctx)
+            if isinstance(iterable_val, CanAssignError):
+                return AnyValue(AnySource.error)
+            return iterable_val
+        else:
+            iterable_val = is_iterable(self.return_annotation, ctx)
+            if isinstance(iterable_val, CanAssignError):
+                return AnyValue(AnySource.error)
+            return iterable_val
 
     def get_generator_send_type(self, ctx: CanAssignContext) -> Value:
         if self.return_annotation is None:
             return AnyValue(AnySource.unannotated)
-        tv_map = get_tv_map(GeneratorValue, self.return_annotation, ctx)
-        if not isinstance(tv_map, CanAssignError):
-            return tv_map.get(SendT, AnyValue(AnySource.generic_argument))
-        # If the return annotation is a non-Generator Iterable, assume the send
-        # type is None.
-        iterable_val = is_iterable(self.return_annotation, ctx)
-        if isinstance(iterable_val, CanAssignError):
-            return AnyValue(AnySource.error)
-        return KnownValue(None)
+        if isinstance(self.node, ast.AsyncFunctionDef):
+            tv_map = get_tv_map(AsyncGeneratorValue, self.return_annotation, ctx)
+            if not isinstance(tv_map, CanAssignError):
+                return tv_map.get(SendT, AnyValue(AnySource.generic_argument))
+            # If the return annotation is a non-Generator Iterable, assume the send
+            # type is None.
+            iterable_val = is_async_iterable(self.return_annotation, ctx)
+            if isinstance(iterable_val, CanAssignError):
+                return AnyValue(AnySource.error)
+            return KnownValue(None)
+        else:
+            tv_map = get_tv_map(GeneratorValue, self.return_annotation, ctx)
+            if not isinstance(tv_map, CanAssignError):
+                return tv_map.get(SendT, AnyValue(AnySource.generic_argument))
+            # If the return annotation is a non-Generator Iterable, assume the send
+            # type is None.
+            iterable_val = is_iterable(self.return_annotation, ctx)
+            if isinstance(iterable_val, CanAssignError):
+                return AnyValue(AnySource.error)
+            return KnownValue(None)
 
     def get_generator_return_type(self, ctx: CanAssignContext) -> Value:
         if self.return_annotation is None:
diff --git a/pyanalyze/name_check_visitor.py b/pyanalyze/name_check_visitor.py
index 13d21d3..c78dad2 100644
--- a/pyanalyze/name_check_visitor.py
+++ b/pyanalyze/name_check_visitor.py
@@ -172,6 +172,7 @@ from .value import (
     GenericBases,
     GenericValue,
     get_tv_map,
+    is_async_iterable,
     is_iterable,
     is_union,
     KnownValue,
@@ -3853,11 +3854,16 @@ class NameCheckVisitor(node_visitor.ReplacingNodeVisitor):
         """
         composite = self.composite_from_node(node)
         if is_async:
-            iterator, _ = self._check_dunder_call(node, composite, "__aiter__", [])
-            anext, _ = self._check_dunder_call(
-                node, Composite(iterator, None, node), "__anext__", []
-            )
-            return self.unpack_awaitable(Composite(anext), node)
+            value = is_async_iterable(composite.value, self)
+            if isinstance(value, CanAssignError):
+                self._show_error_if_checking(
+                    node,
+                    f"{composite.value} is not async iterable",
+                    ErrorCode.unsupported_operation,
+                    detail=str(value),
+                )
+                return AnyValue(AnySource.error)
+            return value
         iterated = composite.value
         result = concrete_values_from_iterable(iterated, self)
         if isinstance(result, CanAssignError):
diff --git a/pyanalyze/value.py b/pyanalyze/value.py
index 4224eb3..6be971b 100644
--- a/pyanalyze/value.py
+++ b/pyanalyze/value.py
@@ -2273,6 +2273,7 @@ def unite_values(*values: Value) -> Value:
 
 T = TypeVar("T")
 IterableValue = GenericValue(collections.abc.Iterable, [TypeVarValue(T)])
+AsyncIterableValue = GenericValue(collections.abc.AsyncIterable, [TypeVarValue(T)])
 
 
 class GetItemProto(Protocol[T]):
@@ -2557,6 +2558,16 @@ def is_iterable(value: Value, ctx: CanAssignContext) -> Union[CanAssignError, Va
     return tv_map.get(T, AnyValue(AnySource.generic_argument))
 
 
+def is_async_iterable(
+    value: Value, ctx: CanAssignContext
+) -> Union[CanAssignError, Value]:
+    """Check whether a value is an async iterable."""
+    tv_map = get_tv_map(AsyncIterableValue, value, ctx)
+    if isinstance(tv_map, CanAssignError):
+        return tv_map
+    return tv_map.get(T, AnyValue(AnySource.generic_argument))
+
+
 def _create_unpacked_list(
     iterable_type: Value, target_length: int, post_starred_length: Optional[int]
 ) -> List[Value]:


Test patch:
diff --git a/pyanalyze/test_async_await.py b/pyanalyze/test_async_await.py
index e6c6e0b..b41343f 100644
--- a/pyanalyze/test_async_await.py
+++ b/pyanalyze/test_async_await.py
@@ -59,7 +59,12 @@ class TestAsyncAwait(TestNameCheckVisitorBase):
 
     @assert_passes()
     def test_async_comprehension(self):
+        from typing_extensions import Self
+
         class ANext:
+            def __aiter__(self) -> Self:
+                return self
+
             async def __anext__(self) -> int:
                 return 42
 
@@ -71,27 +76,6 @@ class TestAsyncAwait(TestNameCheckVisitorBase):
             x = [y async for y in AIter()]
             assert_is_value(x, SequenceValue(list, [(True, TypedValue(int))]))
 
-    @assert_passes()
-    def test_async_generator(self):
-        import collections.abc
-        from typing import AsyncIterator
-
-        async def f() -> AsyncIterator[int]:
-            yield 1
-            yield 2
-
-        async def capybara():
-            x = f()
-            assert_is_value(
-                x, GenericValue(collections.abc.AsyncIterator, [TypedValue(int)])
-            )
-            ints = [i async for i in x]
-            # TODO should be list[int] but we lose the type argument somewhere
-            assert_is_value(
-                ints,
-                SequenceValue(list, [(True, AnyValue(AnySource.generic_argument))]),
-            )
-
     @assert_passes()
     def test_bad_async_comprehension(self):
         async def f():
@@ -293,17 +277,62 @@ class TestNoReturn(TestNameCheckVisitorBase):
 
 class TestAsyncGenerator(TestNameCheckVisitorBase):
     @assert_passes()
-    def test_async_gen(self):
+    def test_async_iterator(self):
         import collections.abc
         from typing import AsyncIterator
 
         async def gen() -> AsyncIterator[int]:
             yield 3
+            yield "not an int"  # E: incompatible_yield
 
         async def capybara() -> None:
             assert_is_value(
                 gen(), GenericValue(collections.abc.AsyncIterator, [TypedValue(int)])
             )
             async for i in gen():
-                # TODO should be int
-                assert_is_value(i, AnyValue(AnySource.generic_argument))
+                assert_is_value(i, TypedValue(int))
+
+    @assert_passes()
+    def test_async_generator(self):
+        import collections.abc
+        from typing import AsyncGenerator
+
+        async def gen() -> AsyncGenerator[int, None]:
+            yield 3
+            yield "not an int"  # E: incompatible_yield
+
+        async def capybara() -> None:
+            assert_is_value(
+                gen(),
+                GenericValue(
+                    collections.abc.AsyncGenerator, [TypedValue(int), KnownValue(None)]
+                ),
+            )
+            async for i in gen():
+                assert_is_value(i, TypedValue(int))
+
+    @assert_passes()
+    def test_async_comprehension_over_generator(self):
+        import collections.abc
+        from typing import AsyncIterator
+
+        async def f() -> AsyncIterator[int]:
+            yield 1
+            yield 2
+
+        async def capybara():
+            x = f()
+            assert_is_value(
+                x, GenericValue(collections.abc.AsyncIterator, [TypedValue(int)])
+            )
+            ints = [i async for i in x]
+            assert_is_value(ints, SequenceValue(list, [(True, TypedValue(int))]))
+
+    @assert_passes()
+    def test_send_type(self):
+        from typing import AsyncGenerator
+
+        async def capybara() -> AsyncGenerator[int, str]:
+            x = yield 1
+            assert_is_value(x, TypedValue(str))
+            yield x  # E: incompatible_yield