[flake8-bugbear] Implement return-in-generator (B901)

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

@@ -0,0 +1,77 @@
+"""
+Should emit:
+B901 - on lines 9, 36
+"""
+
+def broken():
+    if True:
+        return [1, 2, 3]
+
+    yield 3
+    yield 2
+    yield 1
+
+
+def not_broken():
+    if True:
+        return
+
+    yield 3
+    yield 2
+    yield 1
+
+
+def not_broken2():
+    return not_broken()
+
+
+def not_broken3():
+    return
+
+    yield from not_broken()
+
+
+def broken2():
+    return [3, 2, 1]
+
+    yield from not_broken()
+
+
+async def not_broken4():
+    import asyncio
+
+    await asyncio.sleep(1)
+    return 1
+
+
+def not_broken5():
+    def inner():
+        return 2
+
+    yield inner()
+
+
+def not_broken6():
+    return (yield from [])
+
+
+def not_broken7():
+    x = yield from []
+    return x
+
+
+def not_broken8():
+    x = None
+
+    def inner(ex):
+        nonlocal x
+        x = ex
+
+    inner((yield from []))
+    return x
+
+
+class NotBroken9(object):
+    def __await__(self):
+        yield from function()
+        return 42

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

@@ -207,6 +207,9 @@ pub(crate) fn statement(stmt: &Stmt, checker: &mut Checker) {
             if checker.enabled(Rule::MutableArgumentDefault) {
                 flake8_bugbear::rules::mutable_argument_default(checker, function_def);
             }
+            if checker.enabled(Rule::ReturnXInGenerator) {
+                flake8_bugbear::rules::return_x_in_generator(checker, function_def);
+            }
             if checker.any_enabled(&[
                 Rule::UnnecessaryReturnNone,
                 Rule::ImplicitReturnValue,

crates/ruff_linter/src/codes.rs

@@ -383,6 +383,7 @@ pub fn code_to_rule(linter: Linter, code: &str) -> Option<(RuleGroup, Rule)> {
         (Flake8Bugbear, "033") => (RuleGroup::Stable, rules::flake8_bugbear::rules::DuplicateValue),
         (Flake8Bugbear, "034") => (RuleGroup::Stable, rules::flake8_bugbear::rules::ReSubPositionalArgs),
         (Flake8Bugbear, "035") => (RuleGroup::Stable, rules::flake8_bugbear::rules::StaticKeyDictComprehension),
+        (Flake8Bugbear, "901") => (RuleGroup::Preview, rules::flake8_bugbear::rules::ReturnXInGenerator),
         (Flake8Bugbear, "904") => (RuleGroup::Stable, rules::flake8_bugbear::rules::RaiseWithoutFromInsideExcept),
         (Flake8Bugbear, "905") => (RuleGroup::Stable, rules::flake8_bugbear::rules::ZipWithoutExplicitStrict),
         (Flake8Bugbear, "909") => (RuleGroup::Preview, rules::flake8_bugbear::rules::LoopIteratorMutation),

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

@@ -62,6 +62,7 @@ mod tests {
     #[test_case(Rule::UselessContextlibSuppress, Path::new("B022.py"))]
     #[test_case(Rule::UselessExpression, Path::new("B018.ipynb"))]
     #[test_case(Rule::UselessExpression, Path::new("B018.py"))]
+    #[test_case(Rule::ReturnXInGenerator, Path::new("B901.py"))]
     #[test_case(Rule::LoopIteratorMutation, Path::new("B909.py"))]
     fn rules(rule_code: Rule, path: &Path) -> Result<()> {
         let snapshot = format!("{}_{}", rule_code.noqa_code(), path.to_string_lossy());

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

@@ -20,6 +20,7 @@ pub(crate) use raise_literal::*;
 pub(crate) use raise_without_from_inside_except::*;
 pub(crate) use re_sub_positional_args::*;
 pub(crate) use redundant_tuple_in_exception_handler::*;
+pub(crate) use return_x_in_generator::*;
 pub(crate) use reuse_of_groupby_generator::*;
 pub(crate) use setattr_with_constant::*;
 pub(crate) use star_arg_unpacking_after_keyword_arg::*;
@@ -56,6 +57,7 @@ mod raise_literal;
 mod raise_without_from_inside_except;
 mod re_sub_positional_args;
 mod redundant_tuple_in_exception_handler;
+mod return_x_in_generator;
 mod reuse_of_groupby_generator;
 mod setattr_with_constant;
 mod star_arg_unpacking_after_keyword_arg;

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

@@ -0,0 +1,97 @@
+use ruff_diagnostics::Diagnostic;
+use ruff_diagnostics::Violation;
+use ruff_macros::{derive_message_formats, violation};
+use ruff_python_ast::visitor;
+use ruff_python_ast::visitor::Visitor;
+use ruff_python_ast::{self as ast, Expr, Stmt, StmtFunctionDef};
+use ruff_text_size::TextRange;
+
+use crate::checkers::ast::Checker;
+
+/// ## What it does
+/// Checks for `return x` statements in functions, that also contain yield
+/// statements.
+///
+/// ## Why is this bad?
+/// Using `return x` in a generator function used to be syntactically invalid
+/// in Python 2. In Python 3 `return x` can be used in a generator as a return
+/// value in conjunction with yield from. Users coming from Python 2 may expect
+/// the old behavior which might lead to bugs. Use native async def coroutines
+/// or mark intentional return x usage with # noqa on the same line.
+///
+/// ## Example
+/// ```python
+/// def broken():
+///     if True:
+///         return [1, 2, 3]
+///
+///     yield 3
+///     yield 2
+///     yield 1
+/// ```
+#[violation]
+pub struct ReturnXInGenerator;
+
+impl Violation for ReturnXInGenerator {
+    #[derive_message_formats]
+    fn message(&self) -> String {
+        format!("Using `yield` together with `return x`. Use native `async def` coroutines or put a `# noqa` comment on this line if this was intentional.")
+    }
+}
+
+#[derive(Default)]
+struct ReturnXInGeneratorVisitor {
+    return_: Option<TextRange>,
+    has_yield: bool,
+
+    in_expr_statement: bool,
+}
+
+impl Visitor<'_> for ReturnXInGeneratorVisitor {
+    fn visit_stmt(&mut self, stmt: &Stmt) {
+        match stmt {
+            Stmt::FunctionDef(_) => {
+                // do not recurse into nested functions, as they are evaluated
+                // individually
+            }
+            Stmt::Return(ast::StmtReturn { value, range }) => {
+                if Option::is_some(value) {
+                    self.return_ = Some(*range);
+                }
+            }
+            _ => {
+                self.in_expr_statement = stmt.is_expr_stmt();
+                visitor::walk_stmt(self, stmt);
+            }
+        }
+    }
+
+    fn visit_expr(&mut self, expr: &Expr) {
+        if !self.in_expr_statement {
+            return;
+        }
+        match expr {
+            Expr::Yield(_) | Expr::YieldFrom(_) => {
+                self.has_yield = true;
+            }
+            _ => {}
+        }
+    }
+}
+
+/// B901
+pub(crate) fn return_x_in_generator(checker: &mut Checker, function_def: &StmtFunctionDef) {
+    if function_def.name.id == "__await__" {
+        return;
+    }
+
+    let mut visitor = ReturnXInGeneratorVisitor::default();
+    visitor.visit_body(&function_def.body);
+
+    if Option::is_some(&visitor.return_) && visitor.has_yield {
+        checker.diagnostics.push(Diagnostic::new(
+            ReturnXInGenerator,
+            visitor.return_.unwrap(),
+        ));
+    }
+}

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

@@ -0,0 +1,21 @@
+---
+source: crates/ruff_linter/src/rules/flake8_bugbear/mod.rs
+---
+B901.py:8:9: B901 Using `yield` together with `return x`. Use native `async def` coroutines or put a `# noqa` comment on this line if this was intentional.
+   |
+ 6 | def broken():
+ 7 |     if True:
+ 8 |         return [1, 2, 3]
+   |         ^^^^^^^^^^^^^^^^ B901
+ 9 | 
+10 |     yield 3
+   |
+
+B901.py:35:5: B901 Using `yield` together with `return x`. Use native `async def` coroutines or put a `# noqa` comment on this line if this was intentional.
+   |
+34 | def broken2():
+35 |     return [3, 2, 1]
+   |     ^^^^^^^^^^^^^^^^ B901
+36 | 
+37 |     yield from not_broken()
+   |