[ty] Infer yield expression

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

@@ -13,8 +13,7 @@ def inner_generator() -> Generator[int, bytes, str]:
     yield 2
     x = yield 3
 
-    # TODO: this should be `bytes`
-    reveal_type(x)  # revealed: @Todo(yield expressions)
+    reveal_type(x)  # revealed: bytes
 
     return "done"
 
@@ -82,8 +81,7 @@ def inner_generator() -> GeneratorType[int, bytes, str]:
     yield 2
     x = yield 3
 
-    # TODO: this should be `bytes`
-    reveal_type(x)  # revealed: @Todo(yield expressions)
+    reveal_type(x)  # revealed: bytes
 
     return "done"
 
@@ -92,6 +90,88 @@ def outer_generator():
     reveal_type(result)  # revealed: str
 ```
 
+## Infering with type context
+
+A dict literal that is structurally compatible with a `TypedDict` should be accepted.
+
+```py
+from typing import Iterator, TypedDict
+
+class Person(TypedDict):
+    name: str
+
+def persons() -> Iterator[Person]:
+    yield {"name": "Alice"}
+    yield {"name": "Bob"}
+
+    # error: [invalid-yield]
+    # error: [invalid-argument-type]
+    yield {"name": 42}
+```
+
+This also works with `yield from`, where the iterable expression is inferred with the outer
+generator's yield type as type context:
+
+```py
+def persons() -> Iterator[Person]:
+    yield from [{"name": "Alice"}, {"name": "Bob"}]
+
+    # error: [invalid-yield]
+    # error: [invalid-argument-type]
+    yield from [{"name": 42}]
+```
+
+## `yield` expression send type inference
+
+```py
+from typing import AsyncGenerator, AsyncIterator, Generator, Iterator
+
+def unannotated():
+    x = yield 1
+    reveal_type(x)  # revealed: Unknown
+
+def default_generator() -> Generator:
+    x = yield
+    reveal_type(x)  # revealed: None
+
+def generator_one_arg() -> Generator[int]:
+    x = yield 1
+    reveal_type(x)  # revealed: None
+
+def generator_send_str() -> Generator[int, str]:
+    x = yield 1
+    reveal_type(x)  # revealed: str
+
+async def async_generator_default() -> AsyncGenerator[int]:
+    x = yield 1
+    reveal_type(x)  # revealed: None
+
+async def async_generator_send_str() -> AsyncGenerator[int, str]:
+    x = yield 1
+    reveal_type(x)  # revealed: str
+
+def mixing_generator_async_generator() -> Generator[int, int, None] | AsyncGenerator[int, str]:
+    x = yield 1
+    reveal_type(x)  # revealed: int | str
+    return None
+```
+
+`Iterator` has no send type or return type, It is equivalent to using `Generator` with send set to
+`None` and return type to `Unknown`.
+
+```py
+def iterator_send_none() -> Iterator[int]:
+    x = yield 1
+    reveal_type(x)  # revealed: None
+
+async def async_iterator_send_none() -> AsyncIterator[int]:
+    x = yield 1
+    reveal_type(x)  # revealed: None
+
+def iterator_yield_from() -> Generator[int, None, int]:
+    yield from iterator_send_none()
+```
+
 ## Error cases
 
 ### Non-iterable type
@@ -105,12 +185,28 @@ def generator() -> Generator:
 
 ### Invalid `yield` type
 
+<!-- snapshot-diagnostics -->
+
 ```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`
+    # error: [invalid-yield] "Yield type `Literal[""]` does not match annotated yield type `int`"
+    yield ""
+```
+
+### Invalid annotation
+
+```py
+from typing import AsyncGenerator, Generator
+
+def returns_str() -> str:  # error: [invalid-return-type]
+    x = yield 1
+    reveal_type(x)  # revealed: Unknown
+
+def sync_returns_async_generator() -> AsyncGenerator[int, str]:  # error: [invalid-return-type]
+    x = yield 1
+    reveal_type(x)  # revealed: str
 ```
 
 ### Invalid return type
@@ -128,3 +224,31 @@ def invalid_generator2() -> Generator[int, None, None]:
 
     return "done"
 ```
+
+### `yield from` with incompatible yield type
+
+```py
+from typing import Generator
+
+def inner() -> Generator[str, None, None]:
+    yield "hello"
+
+def outer() -> Generator[int, None, None]:
+    # error: [invalid-yield] "Yield type `str` does not match annotated yield type `int`"
+    yield from inner()
+```
+
+### `yield from` with incompatible send type
+
+<!-- snapshot-diagnostics -->
+
+```py
+from typing import Generator
+
+def inner() -> Generator[int, int, None]:
+    x = yield 1
+
+def outer() -> Generator[int, str, None]:
+    # error: [invalid-yield] "Send type `int` does not match annotated send type `str`"
+    yield from inner()
+```

crates/ty_python_semantic/resources/mdtest/snapshots/yield_and_yield_from…_-_`yield`_and_`yield_f…_-_Error_cases_-_Invalid_`yield`_type_(1300c06a97026cce).snap

@@ -0,0 +1,39 @@
+---
+source: crates/ty_test/src/lib.rs
+expression: snapshot
+---
+
+---
+mdtest name: yield_and_yield_from.md - `yield` and `yield from` - Error cases - Invalid `yield` type
+mdtest path: crates/ty_python_semantic/resources/mdtest/expression/yield_and_yield_from.md
+---
+
+# Python source files
+
+## mdtest_snippet.py
+
+```
+1 | from typing import Generator
+2 | 
+3 | def invalid_generator() -> Generator[int, None, None]:
+4 |     # error: [invalid-yield] "Yield type `Literal[""]` does not match annotated yield type `int`"
+5 |     yield ""
+```
+
+# Diagnostics
+
+```
+error[invalid-yield]: Yield expression type does not match annotation
+ --> src/mdtest_snippet.py:3:28
+  |
+1 | from typing import Generator
+2 |
+3 | def invalid_generator() -> Generator[int, None, None]:
+  |                            -------------------------- Function annotated with yield type `int` here
+4 |     # error: [invalid-yield] "Yield type `Literal[""]` does not match annotated yield type `int`"
+5 |     yield ""
+  |           ^^ expression of type `Literal[""]`, expected `int`
+  |
+info: rule `invalid-yield` is enabled by default
+
+```

crates/ty_python_semantic/resources/mdtest/snapshots/yield_and_yield_from…_-_`yield`_and_`yield_f…_-_Error_cases_-_`yield_from`_with_in…_(63388cb3d15fdc10).snap

@@ -0,0 +1,42 @@
+---
+source: crates/ty_test/src/lib.rs
+expression: snapshot
+---
+
+---
+mdtest name: yield_and_yield_from.md - `yield` and `yield from` - Error cases - `yield from` with incompatible send type
+mdtest path: crates/ty_python_semantic/resources/mdtest/expression/yield_and_yield_from.md
+---
+
+# Python source files
+
+## mdtest_snippet.py
+
+```
+1 | from typing import Generator
+2 | 
+3 | def inner() -> Generator[int, int, None]:
+4 |     x = yield 1
+5 | 
+6 | def outer() -> Generator[int, str, None]:
+7 |     # error: [invalid-yield] "Send type `int` does not match annotated send type `str`"
+8 |     yield from inner()
+```
+
+# Diagnostics
+
+```
+error[invalid-yield]: Send type does not match annotation
+ --> src/mdtest_snippet.py:6:16
+  |
+4 |     x = yield 1
+5 |
+6 | def outer() -> Generator[int, str, None]:
+  |                ------------------------- Function annotated with send type `str` here
+7 |     # error: [invalid-yield] "Send type `int` does not match annotated send type `str`"
+8 |     yield from inner()
+  |                ^^^^^^^ generator with send type `int`, expected `str`
+  |
+info: rule `invalid-yield` is enabled by default
+
+```