[flake8-bugbear] Catch yield in subexpressions (B901) (#14453)

crates/ruff_linter/resources/test/fixtures/flake8_bugbear/B901.py

@@ -1,6 +1,6 @@
 """
 Should emit:
-B901 - on lines 9, 36
+B901 - on lines 9, 17, 25
 """
 
 
@@ -13,6 +13,20 @@ def broken():
     yield 1
 
 
+def broken2():
+    return [3, 2, 1]
+
+    yield from not_broken()
+
+
+def broken3(x=None):
+    if x is None:
+        return 0
+
+    for i in range(3):
+        (yield 0) if i > 1 else (yield 1)
+
+
 def not_broken():
     if True:
         return
@@ -32,12 +46,6 @@ def not_broken3():
     yield from not_broken()
 
 
-def broken2():
-    return [3, 2, 1]
-
-    yield from not_broken()
-
-
 async def not_broken4():
     import asyncio
 
@@ -75,4 +83,4 @@ def inner(ex):
 class NotBroken9(object):
     def __await__(self):
         yield from function()
-        return 42
+        return 42

crates/ruff_linter/src/rules/flake8_bugbear/rules/return_in_generator.rs

@@ -1,9 +1,7 @@
 use ruff_diagnostics::Diagnostic;
 use ruff_diagnostics::Violation;
 use ruff_macros::{derive_message_formats, ViolationMetadata};
-use ruff_python_ast::statement_visitor;
-use ruff_python_ast::statement_visitor::StatementVisitor;
-use ruff_python_ast::{self as ast, Expr, Stmt, StmtFunctionDef};
+use ruff_python_ast::{self as ast, visitor::Visitor, Expr, Stmt};
 use ruff_text_size::TextRange;
 
 use crate::checkers::ast::Checker;
@@ -91,7 +89,7 @@ impl Violation for ReturnInGenerator {
 }
 
 /// B901
-pub(crate) fn return_in_generator(checker: &mut Checker, function_def: &StmtFunctionDef) {
+pub(crate) fn return_in_generator(checker: &mut Checker, function_def: &ast::StmtFunctionDef) {
     if function_def.name.id == "__await__" {
         return;
     }
@@ -114,25 +112,37 @@ 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;
                 }
-                _ => {}
+                _ => {
+                    self.visit_expr(value);
+                }
             },
+            Stmt::Assign(_) | Stmt::AnnAssign(_) | Stmt::AugAssign(_) => {}
             Stmt::FunctionDef(_) => {
                 // Do not recurse into nested functions; they're evaluated separately.
             }
-            Stmt::Return(ast::StmtReturn {
-                value: Some(_),
-                range,
-            }) => {
-                self.return_ = Some(*range);
+            Stmt::Return(ast::StmtReturn { value, range }) => {
+                if value.is_some() {
+                    self.return_ = Some(*range);
+                }
+            }
+            _ => ast::visitor::walk_stmt(self, stmt),
+        }
+    }
+
+    fn visit_expr(&mut self, expr: &Expr) {
+        match expr {
+            Expr::Yield(_) | Expr::YieldFrom(_) => {
+                self.has_yield = true;
             }
-            _ => statement_visitor::walk_stmt(self, stmt),
+            Expr::Lambda(_) | Expr::Call(_) => {}
+            _ => ast::visitor::walk_expr(self, expr),
         }
     }
 }

crates/ruff_linter/src/rules/flake8_bugbear/snapshots/ruff_linter__rules__flake8_bugbear__tests__B901_B901.py.snap

@@ -1,5 +1,6 @@
 ---
 source: crates/ruff_linter/src/rules/flake8_bugbear/mod.rs
+snapshot_kind: text
 ---
 B901.py:9:9: B901 Using `yield` and `return {value}` in a generator function can lead to confusing behavior
    |
@@ -11,11 +12,21 @@ B901.py:9:9: B901 Using `yield` and `return {value}` in a generator function can
 11 |     yield 3
    |
 
-B901.py:36:5: B901 Using `yield` and `return {value}` in a generator function can lead to confusing behavior
+B901.py:17:5: B901 Using `yield` and `return {value}` in a generator function can lead to confusing behavior
    |
-35 | def broken2():
-36 |     return [3, 2, 1]
+16 | def broken2():
+17 |     return [3, 2, 1]
    |     ^^^^^^^^^^^^^^^^ B901
-37 |
-38 |     yield from not_broken()
+18 |
+19 |     yield from not_broken()
+   |
+
+B901.py:24:9: B901 Using `yield` and `return {value}` in a generator function can lead to confusing behavior
+   |
+22 | def broken3(x=None):
+23 |     if x is None:
+24 |         return 0
+   |         ^^^^^^^^ B901
+25 |
+26 |     for i in range(3):
    |