Add support for yield from

Author: sharkdp

crates/ty_python_semantic/resources/mdtest/async.md

@@ -6,8 +6,6 @@
 async def retrieve() -> int:
     return 42
 
-reveal_type(retrieve)  # revealed: def retrieve() -> CoroutineType[Any, Any, int]
-
 async def main():
     result = await retrieve()
 
@@ -24,10 +22,102 @@ T = TypeVar("T")
 async def persist(x: T) -> T:
     return x
 
-reveal_type(persist)  # revealed: def persist(x: T) -> CoroutineType[Any, Any, T]
-
 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/expression/yield_and_yield_from.md

@@ -0,0 +1,101 @@
+# `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
+
+        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")
+
+    # The `StopIteration` exception might have a `value` attribute which the default of `None`,
+    # or it could have been customized. So we just return `Unknown` here:
+    reveal_type(result)  # revealed: Unknown
+```
+
+## 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

@@ -32,6 +32,7 @@ class Manager:
 async def test():
     async with Manager() as (x, y):
         reveal_type(x)  # revealed: int
+        reveal_type(y)  # revealed: str
 ```
 
 ## `@asynccontextmanager`

crates/ty_python_semantic/src/types.rs

@@ -4790,35 +4790,48 @@ impl<'db> Type<'db> {
         }
     }
 
+    /// Similar to [`Self::try_enter`], but for async context managers.
+    fn aenter(self, db: &'db dyn Db) -> Type<'db> {
+        // TODO: Add proper error handling and rename this method to `try_aenter`.
+        self.try_call_dunder(db, "__aenter__", CallArguments::none())
+            .map_or(Type::unknown(), |result| {
+                result.return_type(db).resolve_await(db)
+            })
+    }
+
+    /// Resolve the type of an `await …` expression where `self` is the type of the awaitable.
     fn resolve_await(self, db: &'db dyn Db) -> Type<'db> {
+        // TODO: Add proper error handling and rename this method to `try_await`.
         self.try_call_dunder(db, "__await__", CallArguments::none())
             .map_or(Type::unknown(), |result| {
-                let generator_ty = result.return_type(db);
-
-                if let Type::ProtocolInstance(instance) = generator_ty {
-                    if let Protocol::FromClass(class) = instance.inner {
-                        if class.is_known(db, KnownClass::Generator) {
-                            if let Some(specialization) =
-                                class.class_literal_specialized(db, None).1
-                            {
-                                if let [_, _, return_ty] = specialization.types(db) {
-                                    return *return_ty;
-                                }
-                            }
+                result
+                    .return_type(db)
+                    .generator_return_type(db)
+                    .unwrap_or_else(Type::unknown)
+            })
+    }
+
+    /// Get the return type of a `yield from …` expression where `self` is the type of the generator.
+    ///
+    /// This corresponds to the `ReturnT` parameter of the generic `typing.Generator[YieldT, SendT, ReturnT]`
+    /// protocol.
+    fn generator_return_type(self, db: &'db dyn Db) -> Option<Type<'db>> {
+        // TODO: Ideally, we would first try to upcast `self` to an instance of `Generator` and *then*
+        // match on the protocol instance to get the `ReturnType` type parameter.
+
+        if let Type::ProtocolInstance(instance) = self {
+            if let Protocol::FromClass(class) = instance.inner {
+                if class.is_known(db, KnownClass::Generator) {
+                    if let Some(specialization) = class.class_literal_specialized(db, None).1 {
+                        if let [_, _, return_ty] = specialization.types(db) {
+                            return Some(*return_ty);
                         }
                     }
                 }
+            }
+        }
 
-                Type::unknown()
-            })
-    }
-
-    fn aenter(self, db: &'db dyn Db) -> Type<'db> {
-        // TODO: Rename this method to `try_aenter` and add error handling
-        self.try_call_dunder(db, "__aenter__", CallArguments::none())
-            .map_or(Type::unknown(), |result| {
-                result.return_type(db).resolve_await(db)
-            })
+        None
     }
 
     /// Given a class literal or non-dynamic SubclassOf type, try calling it (creating an instance)

crates/ty_python_semantic/src/types/infer.rs

@@ -6092,8 +6092,9 @@ impl<'db, 'ast> TypeInferenceBuilder<'db, 'ast> {
                 err.fallback_element_type(self.db())
             });
 
-        // TODO get type from `ReturnType` of generator
-        todo_type!("Generic `typing.Generator` type")
+        iterable_type
+            .generator_return_type(self.db())
+            .unwrap_or_else(Type::unknown)
     }
 
     fn infer_await_expression(&mut self, await_expression: &ast::ExprAwait) -> Type<'db> {

crates/ty_python_semantic/src/types/signatures.rs

@@ -320,14 +320,14 @@ impl<'db> Signature<'db> {
         inherited_generic_context: Option<GenericContext<'db>>,
         definition: Definition<'db>,
         function_node: &ast::StmtFunctionDef,
-        is_async_generator: bool,
+        is_generator: bool,
     ) -> Self {
         let parameters =
             Parameters::from_parameters(db, definition, function_node.parameters.as_ref());
         let return_ty = function_node.returns.as_ref().map(|returns| {
             let plain_return_ty = definition_expression_type(db, definition, returns.as_ref());
 
-            if function_node.is_async && !is_async_generator {
+            if function_node.is_async && !is_generator {
                 KnownClass::CoroutineType
                     .to_specialized_instance(db, [Type::any(), Type::any(), plain_return_ty])
             } else {

crates/ty_python_semantic/src/types/unpacker.rs

@@ -9,7 +9,7 @@ use crate::Db;
 use crate::semantic_index::ast_ids::node_key::ExpressionNodeKey;
 use crate::semantic_index::scope::ScopeId;
 use crate::types::tuple::{ResizeTupleError, Tuple, TupleLength, TupleSpec, TupleUnpacker};
-use crate::types::{CallArguments, Type, TypeCheckDiagnostics, infer_expression_types};
+use crate::types::{Type, TypeCheckDiagnostics, infer_expression_types};
 use crate::unpack::{UnpackKind, UnpackValue};
 
 use super::context::InferContext;