From 02d3667442e0c8e63989e0276d1e78754cc84b2a Mon Sep 17 00:00:00 2001 From: Danny Yang Date: Sun, 30 Jun 2024 16:15:34 -0400 Subject: [PATCH] Fix typechecking for async generators (#17452) Fixes #10534 This PR fixes a bug in typechecking asynchronous generators. Mypy currently typechecks a generator/comprehension as `AsyncGenerator` if the leftmost expression contains `await`, or if it contains an `async for`. However, there are other situations where we should get async generator: If there is an `await` expression in any of the conditions or in any sequence except for the leftmost one, the generator/comprehension should also be typechecked as `AsyncGenerator`. I've implemented this change in Mypy and added a test case to assert this behavior. If I enter the test cases into a regular repl, I can confirm that the runtime representation is generator/async_generator as the test case expects. According to the [language reference](https://docs.python.org/3/reference/expressions.html#grammar-token-python-grammar-comp_for): > If a comprehension contains either async for clauses or await expressions or other asynchronous comprehensions it is called an asynchronous comprehension. Confusingly, the documentation itself is actually not quite correct either, as pointed out in https://github.com/python/cpython/issues/114104 Alongside this change, I've made a PR to update the docs to be more precise: https://github.com/python/cpython/pull/121175 has more details. --- mypy/checkexpr.py | 9 +++++++-- test-data/unit/check-async-await.test | 19 +++++++++++++++++++ 2 files changed, 26 insertions(+), 2 deletions(-) diff --git a/mypy/checkexpr.py b/mypy/checkexpr.py index fdc0f94b3997..c4ab8a081acc 100644 --- a/mypy/checkexpr.py +++ b/mypy/checkexpr.py @@ -5585,8 +5585,13 @@ def visit_set_comprehension(self, e: SetComprehension) -> Type: def visit_generator_expr(self, e: GeneratorExpr) -> Type: # If any of the comprehensions use async for, the expression will return an async generator - # object, or if the left-side expression uses await. - if any(e.is_async) or has_await_expression(e.left_expr): + # object, or await is used anywhere but in the leftmost sequence. + if ( + any(e.is_async) + or has_await_expression(e.left_expr) + or any(has_await_expression(sequence) for sequence in e.sequences[1:]) + or any(has_await_expression(cond) for condlist in e.condlists for cond in condlist) + ): typ = "typing.AsyncGenerator" # received type is always None in async generator expressions additional_args: list[Type] = [NoneType()] diff --git a/test-data/unit/check-async-await.test b/test-data/unit/check-async-await.test index 876fe0c6be15..0ef08e5a0775 100644 --- a/test-data/unit/check-async-await.test +++ b/test-data/unit/check-async-await.test @@ -573,6 +573,25 @@ async def return_f() -> AsyncGenerator[int, None]: [builtins fixtures/dict.pyi] [typing fixtures/typing-async.pyi] +[case testImplicitAsyncGenerator] +from typing import List + +async def get_list() -> List[int]: + return [1] + +async def predicate() -> bool: + return True + +async def test_implicit_generators() -> None: + reveal_type(await predicate() for _ in [1]) # N: Revealed type is "typing.AsyncGenerator[builtins.bool, None]" + reveal_type(x for x in [1] if await predicate()) # N: Revealed type is "typing.AsyncGenerator[builtins.int, None]" + reveal_type(x for x in await get_list()) # N: Revealed type is "typing.Generator[builtins.int, None, None]" + reveal_type(x for _ in [1] for x in await get_list()) # N: Revealed type is "typing.AsyncGenerator[builtins.int, None]" + +[builtins fixtures/dict.pyi] +[typing fixtures/typing-async.pyi] + + -- The full matrix of coroutine compatibility -- ------------------------------------------