[`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

<!-- How was it tested? -->
I have written tests as directed in #17412

---------

Signed-off-by: 11happy <soni5happy@gmail.com>
Signed-off-by: 11happy <bhuminjaysoni@gmail.com>
Co-authored-by: Brent Westbrook <brentrwestbrook@gmail.com>
Co-authored-by: Brent Westbrook <36778786+ntBre@users.noreply.github.com>

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]

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

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(_)

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);

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),
         }
     }
 }

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]
+   |     ^^^^^^^^^^^^^^^^
+   |

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
+   |     ^^^^^^^^^
+   |

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<TextRange>,
+    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,