[ty] Don't warn yield not in function when yield is in function (#18008)

Author: maxmynter

crates/ruff_linter/resources/test/fixtures/pyflakes/F704.py

@@ -9,3 +9,11 @@ class Foo:
 yield 3
 yield from 3
 await f()
+
+def _():
+    # Invalid yield scopes; but not outside a function
+    type X[T: (yield 1)] = int
+    type Y = (yield 2)
+
+    # Valid yield scope
+    yield 3

crates/ruff_linter/src/checkers/ast/mod.rs

@@ -681,6 +681,17 @@ impl SemanticSyntaxContext for Checker<'_> {
         false
     }
 
+    fn in_yield_allowed_context(&self) -> bool {
+        for scope in self.semantic.current_scopes() {
+            match scope.kind {
+                ScopeKind::Class(_) | ScopeKind::Generator { .. } => return false,
+                ScopeKind::Function(_) | ScopeKind::Lambda(_) => return true,
+                ScopeKind::Module | ScopeKind::Type => {}
+            }
+        }
+        false
+    }
+
     fn in_sync_comprehension(&self) -> bool {
         for scope in self.semantic.current_scopes() {
             if let ScopeKind::Generator {

crates/ruff_linter/src/rules/pyflakes/snapshots/ruff_linter__rules__pyflakes__tests__F704_F704.py.snap

@@ -1,6 +1,5 @@
 ---
 source: crates/ruff_linter/src/rules/pyflakes/mod.rs
-snapshot_kind: text
 ---
 F704.py:6:5: F704 `yield` statement outside of a function
   |
@@ -31,4 +30,6 @@ F704.py:11:1: F704 `await` statement outside of a function
 10 | yield from 3
 11 | await f()
    | ^^^^^^^^^ F704
+12 |
+13 | def _():
    |

crates/ruff_python_parser/src/semantic_errors.rs

@@ -769,16 +769,21 @@ impl SemanticSyntaxChecker {
         // We are intentionally not inspecting the async status of the scope for now to mimic F704.
         // await-outside-async is PLE1142 instead, so we'll end up emitting both syntax errors for
         // cases that trigger F704
+
+        if ctx.in_function_scope() {
+            return;
+        }
+
         if kind.is_await() {
-            if ctx.in_await_allowed_context() {
-                return;
-            }
             // `await` is allowed at the top level of a Jupyter notebook.
             // See: https://ipython.readthedocs.io/en/stable/interactive/autoawait.html.
             if ctx.in_module_scope() && ctx.in_notebook() {
                 return;
             }
-        } else if ctx.in_function_scope() {
+            if ctx.in_await_allowed_context() {
+                return;
+            }
+        } else if ctx.in_yield_allowed_context() {
             return;
         }
 
@@ -1719,6 +1724,35 @@ pub trait SemanticSyntaxContext {
     /// See the trait-level documentation for more details.
     fn in_await_allowed_context(&self) -> bool;
 
+    /// Returns `true` if the visitor is currently in a context where `yield` and `yield from`
+    /// expressions are allowed.
+    ///
+    /// Yield expressions are allowed only in:
+    /// 1. Function definitions
+    /// 2. Lambda expressions
+    ///
+    /// Unlike `await`, yield is not allowed in:
+    /// - Comprehensions (list, set, dict)
+    /// - Generator expressions
+    /// - Class definitions
+    ///
+    /// This method should traverse parent scopes to check if the closest relevant scope
+    /// is a function or lambda, and that no disallowed context (class, comprehension, generator)
+    /// intervenes. For example:
+    ///
+    /// ```python
+    /// def f():
+    ///     yield 1  # okay, in a function
+    ///     lambda: (yield 1)  # okay, in a lambda
+    ///
+    ///     [(yield 1) for x in range(3)]  # error, in a comprehension
+    ///     ((yield 1) for x in range(3))  # error, in a generator expression
+    ///     class C:
+    ///         yield 1  # error, in a class within a function
+    /// ```
+    ///
+    fn in_yield_allowed_context(&self) -> bool;
+
     /// Returns `true` if the visitor is currently inside of a synchronous comprehension.
     ///
     /// This method is necessary because `in_async_context` only checks for the nearest, enclosing

crates/ruff_python_parser/tests/fixtures.rs

@@ -556,6 +556,10 @@ impl SemanticSyntaxContext for SemanticSyntaxCheckerVisitor<'_> {
         true
     }
 
+    fn in_yield_allowed_context(&self) -> bool {
+        true
+    }
+
     fn in_generator_scope(&self) -> bool {
         true
     }

crates/ty_python_semantic/resources/mdtest/diagnostics/semantic_syntax_errors.md

@@ -298,25 +298,25 @@ python-version = "3.12"
 ```
 
 ```py
-# error: [invalid-type-form] "`yield` expressions are not allowed in type expressions"
-# error: [invalid-syntax] "yield expression cannot be used within a TypeVar bound"
-# error: [invalid-syntax] "`yield` statement outside of a function"
-type X[T: (yield 1)] = int
+def _():
+    # error: [invalid-type-form] "`yield` expressions are not allowed in type expressions"
+    # error: [invalid-syntax] "yield expression cannot be used within a TypeVar bound"
+    type X[T: (yield 1)] = int
 
-# error: [invalid-type-form] "`yield` expressions are not allowed in type expressions"
-# error: [invalid-syntax] "yield expression cannot be used within a type alias"
-# error: [invalid-syntax] "`yield` statement outside of a function"
-type Y = (yield 1)
+def _():
+    # error: [invalid-type-form] "`yield` expressions are not allowed in type expressions"
+    # error: [invalid-syntax] "yield expression cannot be used within a type alias"
+    type Y = (yield 1)
 
 # error: [invalid-type-form] "Named expressions are not allowed in type expressions"
 # error: [invalid-syntax] "named expression cannot be used within a generic definition"
 def f[T](x: int) -> (y := 3):
     return x
 
-# error: [invalid-syntax] "`yield from` statement outside of a function"
-# error: [invalid-syntax] "yield expression cannot be used within a generic definition"
-class C[T]((yield from [object])):
-    pass
+def _():
+    # error: [invalid-syntax] "yield expression cannot be used within a generic definition"
+    class C[T]((yield from [object])):
+        pass
 ```
 
 ## `await` outside async function

crates/ty_python_semantic/src/semantic_index/builder.rs

@@ -2481,6 +2481,18 @@ impl SemanticSyntaxContext for SemanticIndexBuilder<'_> {
         false
     }
 
+    fn in_yield_allowed_context(&self) -> bool {
+        for scope_info in self.scope_stack.iter().rev() {
+            let scope = &self.scopes[scope_info.file_scope_id];
+            match scope.kind() {
+                ScopeKind::Class | ScopeKind::Comprehension => return false,
+                ScopeKind::Function | ScopeKind::Lambda => return true,
+                ScopeKind::Module | ScopeKind::TypeAlias | ScopeKind::Annotation => {}
+            }
+        }
+        false
+    }
+
     fn in_sync_comprehension(&self) -> bool {
         for scope_info in self.scope_stack.iter().rev() {
             let scope = &self.scopes[scope_info.file_scope_id];