[ty] Recognise functions containing yield from expressions as being generator functions (#17930)

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

@@ -345,7 +345,7 @@ def f(cond: bool) -> str:
 
 <!-- snapshot-diagnostics -->
 
-A function with a `yield` statement anywhere in its body is a
+A function with a `yield` or `yield from` expression anywhere in its body is a
 [generator function](https://docs.python.org/3/glossary.html#term-generator). A generator function
 implicitly returns an instance of `types.GeneratorType` even if it does not contain any `return`
 statements.
@@ -366,6 +366,9 @@ def h() -> typing.Iterator:
 def i() -> typing.Iterable:
     yield 42
 
+def i2() -> typing.Generator:
+    yield from i()
+
 def j() -> str:  # error: [invalid-return-type]
     yield 42
 ```

crates/ty_python_semantic/resources/mdtest/snapshots/return_type.md_-_Function_return_type_-_Generator_functions.snap

@@ -27,39 +27,42 @@ mdtest path: crates/ty_python_semantic/resources/mdtest/function/return_type.md
 13 | def i() -> typing.Iterable:
 14 |     yield 42
 15 | 
-16 | def j() -> str:  # error: [invalid-return-type]
+16 | def i2() -> typing.Generator:
-17 |     yield 42
+17 |     yield from i()
-18 | import types
+18 | 
-19 | import typing
+19 | def j() -> str:  # error: [invalid-return-type]
-20 | 
+20 |     yield 42
-21 | async def f() -> types.AsyncGeneratorType:
+21 | import types
-22 |     yield 42
+22 | import typing
 23 | 
-24 | async def g() -> typing.AsyncGenerator:
+24 | async def f() -> types.AsyncGeneratorType:
 25 |     yield 42
 26 | 
-27 | async def h() -> typing.AsyncIterator:
+27 | async def g() -> typing.AsyncGenerator:
 28 |     yield 42
 29 | 
-30 | async def i() -> typing.AsyncIterable:
+30 | async def h() -> typing.AsyncIterator:
 31 |     yield 42
 32 | 
-33 | async def j() -> str:  # error: [invalid-return-type]
+33 | async def i() -> typing.AsyncIterable:
 34 |     yield 42
+35 | 
+36 | async def j() -> str:  # error: [invalid-return-type]
+37 |     yield 42
 ```
 
 # Diagnostics
 
 ```
 error: lint:invalid-return-type: Return type does not match returned value
-  --> src/mdtest_snippet.py:16:12
+  --> src/mdtest_snippet.py:19:12
    |
-14 |     yield 42
+17 |     yield from i()
-15 |
+18 |
-16 | def j() -> str:  # error: [invalid-return-type]
+19 | def j() -> str:  # error: [invalid-return-type]
    |            ^^^ Expected `str`, found `types.GeneratorType`
-17 |     yield 42
+20 |     yield 42
-18 | import types
+21 | import types
    |
 info: Function is inferred as returning `types.GeneratorType` because it is a generator function
 info: See https://docs.python.org/3/glossary.html#term-generator for more details
@@ -69,13 +72,13 @@ info: `lint:invalid-return-type` is enabled by default
 
 ```
 error: lint:invalid-return-type: Return type does not match returned value
-  --> src/mdtest_snippet.py:33:18
+  --> src/mdtest_snippet.py:36:18
    |
-31 |     yield 42
-32 |
-33 | async def j() -> str:  # error: [invalid-return-type]
-   |                  ^^^ Expected `str`, found `types.AsyncGeneratorType`
 34 |     yield 42
+35 |
+36 | async def j() -> str:  # error: [invalid-return-type]
+   |                  ^^^ Expected `str`, found `types.AsyncGeneratorType`
+37 |     yield 42
    |
 info: Function is inferred as returning `types.AsyncGeneratorType` because it is an async generator function
 info: See https://docs.python.org/3/glossary.html#term-asynchronous-generator for more details

crates/ty_python_semantic/src/semantic_index/builder.rs

@@ -2310,7 +2310,7 @@ where
 
                 walk_expr(self, expr);
             }
-            ast::Expr::Yield(_) => {
+            ast::Expr::Yield(_) | ast::Expr::YieldFrom(_) => {
                 let scope = self.current_scope();
                 if self.scopes[scope].kind() == ScopeKind::Function {
                     self.generator_functions.insert(scope);