[`pylint`] Fix `PLR1708` false positives on nested functions (#21177)

## Summary

Fixes https://github.com/astral-sh/ruff/issues/21162

## Test Plan

`cargo nextest run pylint`

crates/ruff_linter/resources/test/fixtures/pylint/stop_iteration_return.py

@@ -129,3 +129,26 @@ def generator_with_lambda():
     yield 1
     func = lambda x: x  # Just a regular lambda
     yield 2
+
+# See: https://github.com/astral-sh/ruff/issues/21162
+def foo():
+    def g():
+        yield 1
+    raise StopIteration  # Should not trigger
+
+
+def foo():
+    def g():
+        raise StopIteration  # Should not trigger
+    yield 1
+
+# https://github.com/astral-sh/ruff/pull/21177#pullrequestreview-3430209718
+def foo():
+    yield 1
+    class C:
+        raise StopIteration  # Should trigger
+    yield C
+
+# https://github.com/astral-sh/ruff/pull/21177#discussion_r2539702728
+def foo():
+    raise StopIteration((yield 1))  # Should trigger

crates/ruff_linter/src/checkers/ast/analyze/statement.rs

@@ -131,6 +131,9 @@ pub(crate) fn statement(stmt: &Stmt, checker: &mut Checker) {
             if checker.is_rule_enabled(Rule::GeneratorReturnFromIterMethod) {
                 flake8_pyi::rules::bad_generator_return_type(function_def, checker);
             }
+            if checker.is_rule_enabled(Rule::StopIterationReturn) {
+                pylint::rules::stop_iteration_return(checker, function_def);
+            }
             if checker.source_type.is_stub() {
                 if checker.is_rule_enabled(Rule::StrOrReprDefinedInStub) {
                     flake8_pyi::rules::str_or_repr_defined_in_stub(checker, stmt);
@@ -950,9 +953,6 @@ pub(crate) fn statement(stmt: &Stmt, checker: &mut Checker) {
             if checker.is_rule_enabled(Rule::MisplacedBareRaise) {
                 pylint::rules::misplaced_bare_raise(checker, raise);
             }
-            if checker.is_rule_enabled(Rule::StopIterationReturn) {
-                pylint::rules::stop_iteration_return(checker, raise);
-            }
         }
         Stmt::AugAssign(aug_assign @ ast::StmtAugAssign { target, .. }) => {
             if checker.is_rule_enabled(Rule::GlobalStatement) {

crates/ruff_linter/src/rules/pylint/rules/stop_iteration_return.rs

@@ -1,6 +1,9 @@
 use ruff_macros::{ViolationMetadata, derive_message_formats};
-use ruff_python_ast as ast;
-use ruff_python_ast::visitor::{Visitor, walk_expr, walk_stmt};
+use ruff_python_ast::{
+    self as ast,
+    helpers::map_callable,
+    visitor::{Visitor, walk_expr, walk_stmt},
+};
 use ruff_text_size::Ranged;
 
 use crate::Violation;
@@ -50,65 +53,54 @@ impl Violation for StopIterationReturn {
 }
 
 /// PLR1708
-pub(crate) fn stop_iteration_return(checker: &Checker, raise_stmt: &ast::StmtRaise) {
-    // Fast-path: only continue if this is `raise StopIteration` (with or without args)
-    let Some(exc) = &raise_stmt.exc else {
-        return;
+pub(crate) fn stop_iteration_return(checker: &Checker, function_def: &ast::StmtFunctionDef) {
+    let mut analyzer = GeneratorAnalyzer {
+        checker,
+        has_yield: false,
+        stop_iteration_raises: Vec::new(),
     };
 
-    let is_stop_iteration = match exc.as_ref() {
-        ast::Expr::Call(ast::ExprCall { func, .. }) => {
-            checker.semantic().match_builtin_expr(func, "StopIteration")
+    analyzer.visit_body(&function_def.body);
+
+    if analyzer.has_yield {
+        for raise_stmt in analyzer.stop_iteration_raises {
+            checker.report_diagnostic(StopIterationReturn, raise_stmt.range());
         }
-        expr => checker.semantic().match_builtin_expr(expr, "StopIteration"),
-    };
-
-    if !is_stop_iteration {
-        return;
     }
-
-    // Now check the (more expensive) generator context
-    if !in_generator_context(checker) {
-        return;
-    }
-
-    checker.report_diagnostic(StopIterationReturn, raise_stmt.range());
 }
 
-/// Returns true if we're inside a function that contains any `yield`/`yield from`.
-fn in_generator_context(checker: &Checker) -> bool {
-    for scope in checker.semantic().current_scopes() {
-        if let ruff_python_semantic::ScopeKind::Function(function_def) = scope.kind {
-            if contains_yield_statement(&function_def.body) {
-                return true;
+struct GeneratorAnalyzer<'a, 'b> {
+    checker: &'a Checker<'b>,
+    has_yield: bool,
+    stop_iteration_raises: Vec<&'a ast::StmtRaise>,
+}
+
+impl<'a> Visitor<'a> for GeneratorAnalyzer<'a, '_> {
+    fn visit_stmt(&mut self, stmt: &'a ast::Stmt) {
+        match stmt {
+            ast::Stmt::FunctionDef(_) => {}
+            ast::Stmt::Raise(raise @ ast::StmtRaise { exc: Some(exc), .. }) => {
+                if self
+                    .checker
+                    .semantic()
+                    .match_builtin_expr(map_callable(exc), "StopIteration")
+                {
+                    self.stop_iteration_raises.push(raise);
+                }
+                walk_stmt(self, stmt);
             }
+            _ => walk_stmt(self, stmt),
         }
     }
-    false
-}
 
-/// Check if a statement list contains any yield statements
-fn contains_yield_statement(body: &[ast::Stmt]) -> bool {
-    struct YieldFinder {
-        found: bool,
-    }
-
-    impl Visitor<'_> for YieldFinder {
-        fn visit_expr(&mut self, expr: &ast::Expr) {
-            if matches!(expr, ast::Expr::Yield(_) | ast::Expr::YieldFrom(_)) {
-                self.found = true;
-            } else {
+    fn visit_expr(&mut self, expr: &'a ast::Expr) {
+        match expr {
+            ast::Expr::Lambda(_) => {}
+            ast::Expr::Yield(_) | ast::Expr::YieldFrom(_) => {
+                self.has_yield = true;
                 walk_expr(self, expr);
             }
+            _ => walk_expr(self, expr),
         }
     }
-
-    let mut finder = YieldFinder { found: false };
-    for stmt in body {
-        walk_stmt(&mut finder, stmt);
-        if finder.found {
-            return true;
-        }
-    }
-    false
 }

crates/ruff_linter/src/rules/pylint/snapshots/ruff_linter__rules__pylint__tests__PLR1708_stop_iteration_return.py.snap

@@ -107,3 +107,24 @@ PLR1708 Explicit `raise StopIteration` in generator
     |             ^^^^^^^^^^^^^^^^^^^^^^^^^^^
     |
 help: Use `return` instead
+
+PLR1708 Explicit `raise StopIteration` in generator
+   --> stop_iteration_return.py:149:9
+    |
+147 |     yield 1
+148 |     class C:
+149 |         raise StopIteration  # Should trigger
+    |         ^^^^^^^^^^^^^^^^^^^
+150 |     yield C
+    |
+help: Use `return` instead
+
+PLR1708 Explicit `raise StopIteration` in generator
+   --> stop_iteration_return.py:154:5
+    |
+152 | # https://github.com/astral-sh/ruff/pull/21177#discussion_r2539702728
+153 | def foo():
+154 |     raise StopIteration((yield 1))  # Should trigger
+    |     ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+    |
+help: Use `return` instead