Skip to content

Commit

Permalink
Fix typechecking for async generators (#17452)
Browse files Browse the repository at this point in the history
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
python/cpython#114104

Alongside this change, I've made a PR to update the docs to be more
precise: python/cpython#121175 has more details.
  • Loading branch information
yangdanny97 committed Jun 30, 2024
1 parent 177c8ee commit 02d3667
Show file tree
Hide file tree
Showing 2 changed files with 26 additions and 2 deletions.
9 changes: 7 additions & 2 deletions mypy/checkexpr.py
Original file line number Diff line number Diff line change
Expand Up @@ -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()]
Expand Down
19 changes: 19 additions & 0 deletions test-data/unit/check-async-await.test
Original file line number Diff line number Diff line change
Expand Up @@ -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
-- ------------------------------------------

Expand Down

0 comments on commit 02d3667

Please sign in to comment.