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
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.

Author: yangdanny97

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()]

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
 -- ------------------------------------------