diff --git a/crates/ruff_linter/resources/test/fixtures/pylint/stop_iteration_return.py b/crates/ruff_linter/resources/test/fixtures/pylint/stop_iteration_return.py index a35e31a21e..639cd96845 100644 --- a/crates/ruff_linter/resources/test/fixtures/pylint/stop_iteration_return.py +++ b/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 \ No newline at end of file diff --git a/crates/ruff_linter/src/checkers/ast/analyze/statement.rs b/crates/ruff_linter/src/checkers/ast/analyze/statement.rs index 1232dc52a9..a01194d5be 100644 --- a/crates/ruff_linter/src/checkers/ast/analyze/statement.rs +++ b/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) { diff --git a/crates/ruff_linter/src/rules/pylint/rules/stop_iteration_return.rs b/crates/ruff_linter/src/rules/pylint/rules/stop_iteration_return.rs index 2abc339592..6ae47fe66b 100644 --- a/crates/ruff_linter/src/rules/pylint/rules/stop_iteration_return.rs +++ b/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 } diff --git a/crates/ruff_linter/src/rules/pylint/snapshots/ruff_linter__rules__pylint__tests__PLR1708_stop_iteration_return.py.snap b/crates/ruff_linter/src/rules/pylint/snapshots/ruff_linter__rules__pylint__tests__PLR1708_stop_iteration_return.py.snap index 5773ff81ea..a6775f05d7 100644 --- a/crates/ruff_linter/src/rules/pylint/snapshots/ruff_linter__rules__pylint__tests__PLR1708_stop_iteration_return.py.snap +++ b/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