[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 astral-sh/ty#151

closes astral-sh/ty#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](https://github.com/PrefectHQ/prefect/blob/e926b8c4c114e74533e7bd75d83e7f205c774645/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](https://github.com/mitmproxy/mitmproxy/blob/a4d794c59a27472d193a592d8037505a1cf6ae93/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

Author: sharkdp

crates/ty_python_semantic/resources/mdtest/async.md

@@ -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
+```

crates/ty_python_semantic/resources/mdtest/call/function.md

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

crates/ty_python_semantic/resources/mdtest/expression/yield_and_yield_from.md

@@ -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"
+```

crates/ty_python_semantic/resources/mdtest/with/async.md

@@ -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())
 ```

crates/ty_python_semantic/src/semantic_index/builder.rs

@@ -2805,7 +2805,9 @@ impl<'ast> Unpackable<'ast> {
         match self {
             Unpackable::Assign(_) => UnpackKind::Assign,
             Unpackable::For(_) | Unpackable::Comprehension { .. } => UnpackKind::Iterable,
-            Unpackable::WithItem { .. } => UnpackKind::ContextManager,
+            Unpackable::WithItem { is_async, .. } => UnpackKind::ContextManager {
+                is_async: *is_async,
+            },
         }
     }