From c722f498fe409f2897c746d0f40397a9460496e4 Mon Sep 17 00:00:00 2001 From: Bhuminjay Soni Date: Wed, 3 Dec 2025 22:35:15 +0530 Subject: [PATCH] [`flake8-bugbear`] Catch `yield` expressions within other statements (`B901`) (#21200) ## Summary This PR re-implements [return-in-generator (B901)](https://docs.astral.sh/ruff/rules/return-in-generator/#return-in-generator-b901) for async generators as a semantic syntax error. This is not a syntax error for sync generators, so we'll need to preserve both the lint rule and the syntax error in this case. It also updates B901 and the new implementation to catch cases where the generator's `yield` or `yield from` expression is part of another statement, as in: ```py def foo(): return (yield) ``` These were previously not caught because we only looked for `Stmt::Expr(Expr::Yield)` in `visit_stmt` instead of visiting `yield` expressions directly. I think this modification is within the spirit of the rule and safe to try out since the rule is in preview. ## Test Plan I have written tests as directed in #17412 --------- Signed-off-by: 11happy Signed-off-by: 11happy Co-authored-by: Brent Westbrook Co-authored-by: Brent Westbrook <36778786+ntBre@users.noreply.github.com> --- .../test/fixtures/flake8_bugbear/B901.py | 16 ++++- .../syntax_errors/return_in_generator.py | 24 ++++++++ crates/ruff_linter/src/checkers/ast/mod.rs | 7 +++ crates/ruff_linter/src/linter.rs | 1 + .../rules/return_in_generator.rs | 29 ++++++--- ...__flake8_bugbear__tests__B901_B901.py.snap | 43 +++++++++++++ ...linter__tests__return_in_generator.py.snap | 55 +++++++++++++++++ .../ruff_python_parser/src/semantic_errors.rs | 61 ++++++++++++++++++- 8 files changed, 220 insertions(+), 16 deletions(-) create mode 100644 crates/ruff_linter/resources/test/fixtures/syntax_errors/return_in_generator.py create mode 100644 crates/ruff_linter/src/snapshots/ruff_linter__linter__tests__return_in_generator.py.snap diff --git a/crates/ruff_linter/resources/test/fixtures/flake8_bugbear/B901.py b/crates/ruff_linter/resources/test/fixtures/flake8_bugbear/B901.py index 42fdda60d7..acb932f25f 100644 --- a/crates/ruff_linter/resources/test/fixtures/flake8_bugbear/B901.py +++ b/crates/ruff_linter/resources/test/fixtures/flake8_bugbear/B901.py @@ -52,16 +52,16 @@ def not_broken5(): yield inner() -def not_broken6(): +def broken3(): return (yield from []) -def not_broken7(): +def broken4(): x = yield from [] return x -def not_broken8(): +def broken5(): x = None def inner(ex): @@ -76,3 +76,13 @@ class NotBroken9(object): def __await__(self): yield from function() return 42 + + +async def broken6(): + yield 1 + return foo() + + +async def broken7(): + yield 1 + return [1, 2, 3] diff --git a/crates/ruff_linter/resources/test/fixtures/syntax_errors/return_in_generator.py b/crates/ruff_linter/resources/test/fixtures/syntax_errors/return_in_generator.py new file mode 100644 index 0000000000..56ef47c17b --- /dev/null +++ b/crates/ruff_linter/resources/test/fixtures/syntax_errors/return_in_generator.py @@ -0,0 +1,24 @@ +async def gen(): + yield 1 + return 42 + +def gen(): # B901 but not a syntax error - not an async generator + yield 1 + return 42 + +async def gen(): # ok - no value in return + yield 1 + return + +async def gen(): + yield 1 + return foo() + +async def gen(): + yield 1 + return [1, 2, 3] + +async def gen(): + if True: + yield 1 + return 10 diff --git a/crates/ruff_linter/src/checkers/ast/mod.rs b/crates/ruff_linter/src/checkers/ast/mod.rs index aa9fa839f2..4d4d7e9293 100644 --- a/crates/ruff_linter/src/checkers/ast/mod.rs +++ b/crates/ruff_linter/src/checkers/ast/mod.rs @@ -69,6 +69,7 @@ use crate::noqa::NoqaMapping; use crate::package::PackageRoot; use crate::preview::is_undefined_export_in_dunder_init_enabled; use crate::registry::Rule; +use crate::rules::flake8_bugbear::rules::ReturnInGenerator; use crate::rules::pyflakes::rules::{ LateFutureImport, MultipleStarredExpressions, ReturnOutsideFunction, UndefinedLocalWithNestedImportStarUsage, YieldOutsideFunction, @@ -729,6 +730,12 @@ impl SemanticSyntaxContext for Checker<'_> { self.report_diagnostic(NonlocalWithoutBinding { name }, error.range); } } + SemanticSyntaxErrorKind::ReturnInGenerator => { + // B901 + if self.is_rule_enabled(Rule::ReturnInGenerator) { + self.report_diagnostic(ReturnInGenerator, error.range); + } + } SemanticSyntaxErrorKind::ReboundComprehensionVariable | SemanticSyntaxErrorKind::DuplicateTypeParameter | SemanticSyntaxErrorKind::MultipleCaseAssignment(_) diff --git a/crates/ruff_linter/src/linter.rs b/crates/ruff_linter/src/linter.rs index 3ec070dd26..719d5ac9c5 100644 --- a/crates/ruff_linter/src/linter.rs +++ b/crates/ruff_linter/src/linter.rs @@ -1043,6 +1043,7 @@ mod tests { Rule::YieldFromInAsyncFunction, Path::new("yield_from_in_async_function.py") )] + #[test_case(Rule::ReturnInGenerator, Path::new("return_in_generator.py"))] fn test_syntax_errors(rule: Rule, path: &Path) -> Result<()> { let snapshot = path.to_string_lossy().to_string(); let path = Path::new("resources/test/fixtures/syntax_errors").join(path); diff --git a/crates/ruff_linter/src/rules/flake8_bugbear/rules/return_in_generator.rs b/crates/ruff_linter/src/rules/flake8_bugbear/rules/return_in_generator.rs index f7584dd4bb..0b089b3459 100644 --- a/crates/ruff_linter/src/rules/flake8_bugbear/rules/return_in_generator.rs +++ b/crates/ruff_linter/src/rules/flake8_bugbear/rules/return_in_generator.rs @@ -1,6 +1,5 @@ use ruff_macros::{ViolationMetadata, derive_message_formats}; -use ruff_python_ast::statement_visitor; -use ruff_python_ast::statement_visitor::StatementVisitor; +use ruff_python_ast::visitor::{Visitor, walk_expr, walk_stmt}; use ruff_python_ast::{self as ast, Expr, Stmt, StmtFunctionDef}; use ruff_text_size::TextRange; @@ -96,6 +95,11 @@ pub(crate) fn return_in_generator(checker: &Checker, function_def: &StmtFunction return; } + // Async functions are flagged by the `ReturnInGenerator` semantic syntax error. + if function_def.is_async { + return; + } + let mut visitor = ReturnInGeneratorVisitor::default(); visitor.visit_body(&function_def.body); @@ -112,15 +116,9 @@ struct ReturnInGeneratorVisitor { has_yield: bool, } -impl StatementVisitor<'_> for ReturnInGeneratorVisitor { +impl Visitor<'_> for ReturnInGeneratorVisitor { fn visit_stmt(&mut self, stmt: &Stmt) { match stmt { - Stmt::Expr(ast::StmtExpr { value, .. }) => match **value { - Expr::Yield(_) | Expr::YieldFrom(_) => { - self.has_yield = true; - } - _ => {} - }, Stmt::FunctionDef(_) => { // Do not recurse into nested functions; they're evaluated separately. } @@ -130,8 +128,19 @@ impl StatementVisitor<'_> for ReturnInGeneratorVisitor { node_index: _, }) => { self.return_ = Some(*range); + walk_stmt(self, stmt); } - _ => statement_visitor::walk_stmt(self, stmt), + _ => walk_stmt(self, stmt), + } + } + + fn visit_expr(&mut self, expr: &Expr) { + match expr { + Expr::Lambda(_) => {} + Expr::Yield(_) | Expr::YieldFrom(_) => { + self.has_yield = true; + } + _ => walk_expr(self, expr), } } } diff --git a/crates/ruff_linter/src/rules/flake8_bugbear/snapshots/ruff_linter__rules__flake8_bugbear__tests__B901_B901.py.snap b/crates/ruff_linter/src/rules/flake8_bugbear/snapshots/ruff_linter__rules__flake8_bugbear__tests__B901_B901.py.snap index 951860f81e..21bf1b1645 100644 --- a/crates/ruff_linter/src/rules/flake8_bugbear/snapshots/ruff_linter__rules__flake8_bugbear__tests__B901_B901.py.snap +++ b/crates/ruff_linter/src/rules/flake8_bugbear/snapshots/ruff_linter__rules__flake8_bugbear__tests__B901_B901.py.snap @@ -21,3 +21,46 @@ B901 Using `yield` and `return {value}` in a generator function can lead to conf 37 | 38 | yield from not_broken() | + +B901 Using `yield` and `return {value}` in a generator function can lead to confusing behavior + --> B901.py:56:5 + | +55 | def broken3(): +56 | return (yield from []) + | ^^^^^^^^^^^^^^^^^^^^^^ + | + +B901 Using `yield` and `return {value}` in a generator function can lead to confusing behavior + --> B901.py:61:5 + | +59 | def broken4(): +60 | x = yield from [] +61 | return x + | ^^^^^^^^ + | + +B901 Using `yield` and `return {value}` in a generator function can lead to confusing behavior + --> B901.py:72:5 + | +71 | inner((yield from [])) +72 | return x + | ^^^^^^^^ + | + +B901 Using `yield` and `return {value}` in a generator function can lead to confusing behavior + --> B901.py:83:5 + | +81 | async def broken6(): +82 | yield 1 +83 | return foo() + | ^^^^^^^^^^^^ + | + +B901 Using `yield` and `return {value}` in a generator function can lead to confusing behavior + --> B901.py:88:5 + | +86 | async def broken7(): +87 | yield 1 +88 | return [1, 2, 3] + | ^^^^^^^^^^^^^^^^ + | diff --git a/crates/ruff_linter/src/snapshots/ruff_linter__linter__tests__return_in_generator.py.snap b/crates/ruff_linter/src/snapshots/ruff_linter__linter__tests__return_in_generator.py.snap new file mode 100644 index 0000000000..2abd1fad09 --- /dev/null +++ b/crates/ruff_linter/src/snapshots/ruff_linter__linter__tests__return_in_generator.py.snap @@ -0,0 +1,55 @@ +--- +source: crates/ruff_linter/src/linter.rs +--- +B901 Using `yield` and `return {value}` in a generator function can lead to confusing behavior + --> resources/test/fixtures/syntax_errors/return_in_generator.py:3:5 + | +1 | async def gen(): +2 | yield 1 +3 | return 42 + | ^^^^^^^^^ +4 | +5 | def gen(): # B901 but not a syntax error - not an async generator + | + +B901 Using `yield` and `return {value}` in a generator function can lead to confusing behavior + --> resources/test/fixtures/syntax_errors/return_in_generator.py:7:5 + | +5 | def gen(): # B901 but not a syntax error - not an async generator +6 | yield 1 +7 | return 42 + | ^^^^^^^^^ +8 | +9 | async def gen(): # ok - no value in return + | + +B901 Using `yield` and `return {value}` in a generator function can lead to confusing behavior + --> resources/test/fixtures/syntax_errors/return_in_generator.py:15:5 + | +13 | async def gen(): +14 | yield 1 +15 | return foo() + | ^^^^^^^^^^^^ +16 | +17 | async def gen(): + | + +B901 Using `yield` and `return {value}` in a generator function can lead to confusing behavior + --> resources/test/fixtures/syntax_errors/return_in_generator.py:19:5 + | +17 | async def gen(): +18 | yield 1 +19 | return [1, 2, 3] + | ^^^^^^^^^^^^^^^^ +20 | +21 | async def gen(): + | + +B901 Using `yield` and `return {value}` in a generator function can lead to confusing behavior + --> resources/test/fixtures/syntax_errors/return_in_generator.py:24:5 + | +22 | if True: +23 | yield 1 +24 | return 10 + | ^^^^^^^^^ + | diff --git a/crates/ruff_python_parser/src/semantic_errors.rs b/crates/ruff_python_parser/src/semantic_errors.rs index 2c573271e1..0c7ceef4a4 100644 --- a/crates/ruff_python_parser/src/semantic_errors.rs +++ b/crates/ruff_python_parser/src/semantic_errors.rs @@ -3,12 +3,13 @@ //! This checker is not responsible for traversing the AST itself. Instead, its //! [`SemanticSyntaxChecker::visit_stmt`] and [`SemanticSyntaxChecker::visit_expr`] methods should //! be called in a parent `Visitor`'s `visit_stmt` and `visit_expr` methods, respectively. + use ruff_python_ast::{ self as ast, Expr, ExprContext, IrrefutablePatternKind, Pattern, PythonVersion, Stmt, StmtExpr, - StmtImportFrom, + StmtFunctionDef, StmtImportFrom, comparable::ComparableExpr, helpers, - visitor::{Visitor, walk_expr}, + visitor::{Visitor, walk_expr, walk_stmt}, }; use ruff_text_size::{Ranged, TextRange, TextSize}; use rustc_hash::{FxBuildHasher, FxHashSet}; @@ -739,7 +740,21 @@ impl SemanticSyntaxChecker { self.seen_futures_boundary = true; } } - Stmt::FunctionDef(_) => { + Stmt::FunctionDef(StmtFunctionDef { is_async, body, .. }) => { + if *is_async { + let mut visitor = ReturnVisitor::default(); + visitor.visit_body(body); + + if visitor.has_yield { + if let Some(return_range) = visitor.return_range { + Self::add_error( + ctx, + SemanticSyntaxErrorKind::ReturnInGenerator, + return_range, + ); + } + } + } self.seen_futures_boundary = true; } _ => { @@ -1213,6 +1228,9 @@ impl Display for SemanticSyntaxError { SemanticSyntaxErrorKind::NonlocalWithoutBinding(name) => { write!(f, "no binding for nonlocal `{name}` found") } + SemanticSyntaxErrorKind::ReturnInGenerator => { + write!(f, "`return` with value in async generator") + } } } } @@ -1619,6 +1637,9 @@ pub enum SemanticSyntaxErrorKind { /// Represents a default type parameter followed by a non-default type parameter. TypeParameterDefaultOrder(String), + + /// Represents a `return` statement with a value in an asynchronous generator. + ReturnInGenerator, } #[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, get_size2::GetSize)] @@ -1735,6 +1756,40 @@ impl Visitor<'_> for ReboundComprehensionVisitor<'_> { } } +#[derive(Default)] +struct ReturnVisitor { + return_range: Option, + has_yield: bool, +} + +impl Visitor<'_> for ReturnVisitor { + fn visit_stmt(&mut self, stmt: &Stmt) { + match stmt { + // Do not recurse into nested functions; they're evaluated separately. + Stmt::FunctionDef(_) | Stmt::ClassDef(_) => {} + Stmt::Return(ast::StmtReturn { + value: Some(_), + range, + .. + }) => { + self.return_range = Some(*range); + walk_stmt(self, stmt); + } + _ => walk_stmt(self, stmt), + } + } + + fn visit_expr(&mut self, expr: &Expr) { + match expr { + Expr::Lambda(_) => {} + Expr::Yield(_) | Expr::YieldFrom(_) => { + self.has_yield = true; + } + _ => walk_expr(self, expr), + } + } +} + struct MatchPatternVisitor<'a, Ctx> { names: FxHashSet<&'a ast::name::Name>, ctx: &'a Ctx,