[ty] Support async/await, async with and yield from (#19595)

## Summary

- Add support for the return types of `async` functions
- Add type inference for `await` expressions
- Add support for `async with` / async context managers
- Add support for `yield from` expressions

This PR is generally lacking proper error handling in some cases (e.g.
illegal `__await__` attributes). I'm planning to work on this in a
follow-up.

part of https://github.com/astral-sh/ty/issues/151

closes https://github.com/astral-sh/ty/issues/736

## Ecosystem

There are a lot of true positives on `prefect` which look similar to:
```diff
prefect (https://github.com/PrefectHQ/prefect)
+ src/integrations/prefect-aws/tests/workers/test_ecs_worker.py:406:12: error[unresolved-attribute] Type `str` has no attribute `status_code`
```

This is due to a wrong return type annotation
[here](e926b8c4c1/src/integrations/prefect-aws/tests/workers/test_ecs_worker.py (L355-L391)).

```diff
mitmproxy (https://github.com/mitmproxy/mitmproxy)
+ test/mitmproxy/addons/test_clientplayback.py:18:1: error[invalid-argument-type] Argument to function `asynccontextmanager` is incorrect: Expected `(...) -> AsyncIterator[Unknown]`, found `def tcp_server(handle_conn, **server_args) -> Unknown | tuple[str, int]`
```


[This](a4d794c59a/test/mitmproxy/addons/test_clientplayback.py (L18-L19))
is a true positive. That function should return
`AsyncIterator[Address]`, not `Address`.

I looked through almost all of the other new diagnostics and they all
look like known problems or true positives.

## Typing conformance

The typing conformance diff looks good.

## Test Plan

New Markdown tests
This commit is contained in:
David Peter 2025-07-30 11:51:21 +02:00 committed by GitHub
parent c5ac998892
commit 4ecf1d205a
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
12 changed files with 472 additions and 46 deletions

View file

@ -0,0 +1,123 @@
# `async` / `await`
## Basic
```py
async def retrieve() -> int:
return 42
async def main():
result = await retrieve()
reveal_type(result) # revealed: int
```
## Generic `async` functions
```py
from typing import TypeVar
T = TypeVar("T")
async def persist(x: T) -> T:
return x
async def f(x: int):
result = await persist(x)
reveal_type(result) # revealed: int
```
## Use cases
### `Future`
```py
import asyncio
import concurrent.futures
def blocking_function() -> int:
return 42
async def main():
loop = asyncio.get_event_loop()
with concurrent.futures.ThreadPoolExecutor() as pool:
result = await loop.run_in_executor(pool, blocking_function)
# TODO: should be `int`
reveal_type(result) # revealed: Unknown
```
### `asyncio.Task`
```py
import asyncio
async def f() -> int:
return 1
async def main():
task = asyncio.create_task(f())
result = await task
# TODO: this should be `int`
reveal_type(result) # revealed: Unknown
```
### `asyncio.gather`
```py
import asyncio
async def task(name: str) -> int:
return len(name)
async def main():
(a, b) = await asyncio.gather(
task("A"),
task("B"),
)
# TODO: these should be `int`
reveal_type(a) # revealed: Unknown
reveal_type(b) # revealed: Unknown
```
## Under the hood
```toml
[environment]
python-version = "3.12" # Use 3.12 to be able to use PEP 695 generics
```
Let's look at the example from the beginning again:
```py
async def retrieve() -> int:
return 42
```
When we look at the signature of this function, we see that it actually returns a `CoroutineType`:
```py
reveal_type(retrieve) # revealed: def retrieve() -> CoroutineType[Any, Any, int]
```
The expression `await retrieve()` desugars into a call to the `__await__` dunder method on the
`CoroutineType` object, followed by a `yield from`. Let's first see the return type of `__await__`:
```py
reveal_type(retrieve().__await__()) # revealed: Generator[Any, None, int]
```
We can see that this returns a `Generator` that yields `Any`, and eventually returns `int`. For the
final type of the `await` expression, we retrieve that third argument of the `Generator` type:
```py
from typing import Generator
def _():
result = yield from retrieve().__await__()
reveal_type(result) # revealed: int
```

View file

@ -15,8 +15,7 @@ reveal_type(get_int()) # revealed: int
async def get_int_async() -> int:
return 42
# TODO: we don't yet support `types.CoroutineType`, should be generic `Coroutine[Any, Any, int]`
reveal_type(get_int_async()) # revealed: @Todo(generic types.CoroutineType)
reveal_type(get_int_async()) # revealed: CoroutineType[Any, Any, int]
```
## Generic

View file

@ -0,0 +1,130 @@
# `yield` and `yield from`
## Basic `yield` and `yield from`
The type of a `yield` expression is the "send" type of the generator function. The type of a
`yield from` expression is the return type of the inner generator:
```py
from typing import Generator
def inner_generator() -> Generator[int, bytes, str]:
yield 1
yield 2
x = yield 3
# TODO: this should be `bytes`
reveal_type(x) # revealed: @Todo(yield expressions)
return "done"
def outer_generator():
result = yield from inner_generator()
reveal_type(result) # revealed: str
```
## `yield from` with a custom iterable
`yield from` can also be used with custom iterable types. In that case, the type of the `yield from`
expression can not be determined
```py
from typing import Generator, TypeVar, Generic
T = TypeVar("T")
class OnceIterator(Generic[T]):
def __init__(self, value: T):
self.value = value
self.returned = False
def __next__(self) -> T:
if self.returned:
raise StopIteration(42)
self.returned = True
return self.value
class Once(Generic[T]):
def __init__(self, value: T):
self.value = value
def __iter__(self) -> OnceIterator[T]:
return OnceIterator(self.value)
for x in Once("a"):
reveal_type(x) # revealed: str
def generator() -> Generator:
result = yield from Once("a")
# At runtime, the value of `result` will be the `.value` attribute of the `StopIteration`
# error raised by `OnceIterator` to signal to the interpreter that the iterator has been
# exhausted. Here that will always be 42, but this information cannot be captured in the
# signature of `OnceIterator.__next__`, since exceptions lie outside the type signature.
# We therefore just infer `Unknown` here.
#
# If the `StopIteration` error in `OnceIterator.__next__` had been simply `raise StopIteration`
# (the more common case), then the `.value` attribute of the `StopIteration` instance
# would default to `None`.
reveal_type(result) # revealed: Unknown
```
## `yield from` with a generator that return `types.GeneratorType`
`types.GeneratorType` is a nominal type that implements the `typing.Generator` protocol:
```py
from types import GeneratorType
def inner_generator() -> GeneratorType[int, bytes, str]:
yield 1
yield 2
x = yield 3
# TODO: this should be `bytes`
reveal_type(x) # revealed: @Todo(yield expressions)
return "done"
def outer_generator():
result = yield from inner_generator()
reveal_type(result) # revealed: str
```
## Error cases
### Non-iterable type
```py
from typing import Generator
def generator() -> Generator:
yield from 42 # error: [not-iterable] "Object of type `Literal[42]` is not iterable"
```
### Invalid `yield` type
```py
from typing import Generator
# TODO: This should be an error. Claims to yield `int`, but yields `str`.
def invalid_generator() -> Generator[int, None, None]:
yield "not an int" # This should be an `int`
```
### Invalid return type
```py
from typing import Generator
# TODO: should emit an error (does not return `str`)
def invalid_generator1() -> Generator[int, None, str]:
yield 1
# TODO: should emit an error (does not return `int`)
def invalid_generator2() -> Generator[int, None, None]:
yield 1
return "done"
```

View file

@ -17,5 +17,80 @@ class Manager:
async def test():
async with Manager() as f:
reveal_type(f) # revealed: @Todo(async `with` statement)
reveal_type(f) # revealed: Target
```
## Multiple targets
```py
class Manager:
async def __aenter__(self) -> tuple[int, str]:
return 42, "hello"
async def __aexit__(self, exc_type, exc_value, traceback): ...
async def test():
async with Manager() as (x, y):
reveal_type(x) # revealed: int
reveal_type(y) # revealed: str
```
## `@asynccontextmanager`
```py
from contextlib import asynccontextmanager
from typing import AsyncGenerator
class Session: ...
@asynccontextmanager
async def connect() -> AsyncGenerator[Session]:
yield Session()
# TODO: this should be `() -> _AsyncGeneratorContextManager[Session, None]`
reveal_type(connect) # revealed: (...) -> _AsyncGeneratorContextManager[Unknown, None]
async def main():
async with connect() as session:
# TODO: should be `Session`
reveal_type(session) # revealed: Unknown
```
## `asyncio.timeout`
```toml
[environment]
python-version = "3.11"
```
```py
import asyncio
async def long_running_task():
await asyncio.sleep(5)
async def main():
async with asyncio.timeout(1):
await long_running_task()
```
## `asyncio.TaskGroup`
```toml
[environment]
python-version = "3.11"
```
```py
import asyncio
async def long_running_task():
await asyncio.sleep(5)
async def main():
async with asyncio.TaskGroup() as tg:
# TODO: should be `TaskGroup`
reveal_type(tg) # revealed: Unknown
tg.create_task(long_running_task())
```