refactor: use visit_stmt & have siilar impl to original B901

Signed-off-by: 11happy <soni5happy@gmail.com>
astral-sh · ntBre · Dec 3, 2025 · Nov 2, 2025 · Nov 7, 2025 · Nov 12, 2025
commit bd766dc7f51db4fb101a6bf65732f33db84a33d4
diff --git a/crates/ruff_linter/resources/test/fixtures/semantic_errors/return_in_async_generator.py b/crates/ruff_linter/resources/test/fixtures/semantic_errors/return_in_async_generator.py
@@ -1,3 +1,24 @@
 async def gen():
     yield 1
-    return 42
+    return 42
+
+def gen(): #ok
+    yield 1
+    return 42
+
+async def gen(): #ok
+    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
@@ -34,8 +34,6 @@ use ruff_notebook::{CellOffsets, NotebookIndex};
 use ruff_python_ast::helpers::{collect_import_from_member, is_docstring_stmt, to_module_path};
 use ruff_python_ast::identifier::Identifier;
 use ruff_python_ast::name::QualifiedName;
-use ruff_python_ast::statement_visitor;
-use ruff_python_ast::statement_visitor::StatementVisitor;
 use ruff_python_ast::str::Quote;
 use ruff_python_ast::visitor::{Visitor, walk_except_handler, walk_pattern};
 use ruff_python_ast::{
@@ -619,29 +617,6 @@ pub(crate) struct TypingImporter<'a, 'b> {
     member: &'a str,
 }
 
-#[derive(Default)]
-struct ReturnVisitor {
-    return_range: Option<TextRange>,
-}
-
-impl StatementVisitor<'_> for ReturnVisitor {
-    fn visit_stmt(&mut self, stmt: &Stmt) {
-        match stmt {
-            Stmt::FunctionDef(_) | Stmt::ClassDef(_) => {}
-            Stmt::Return(ast::StmtReturn {
-                value: Some(_),
-                range,
-                ..
-            }) => {
-                if self.return_range.is_none() {
-                    self.return_range = Some(*range);
-                }
-            }
-            _ => statement_visitor::walk_stmt(self, stmt),
-        }
-    }
-}
-
 impl TypingImporter<'_, '_> {
     /// Create an [`Edit`] that makes the requested symbol available at `position`.
     ///
@@ -864,24 +839,6 @@ impl SemanticSyntaxContext for Checker<'_> {
         )
     }
 
-    fn has_return(&self) -> Option<TextRange> {
-        for scope in self.semantic.current_scopes() {
-            match &scope.kind {
-                ScopeKind::Function(ast::StmtFunctionDef { body, .. }) => {
-                    let mut visitor = ReturnVisitor::default();
-                    visitor.visit_body(body);
-                    return visitor.return_range;
-                }
-                ScopeKind::Lambda(_) | ScopeKind::Class(_) => return None,
-                ScopeKind::Generator { .. }
-                | ScopeKind::Module
-                | ScopeKind::Type
-                | ScopeKind::DunderClassCell => {}
-            }
-        }
-        None
-    }
-
     fn in_loop_context(&self) -> bool {
         let mut child = self.semantic.current_statement();
 

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
@@ -96,6 +96,10 @@ pub(crate) fn return_in_generator(checker: &Checker, function_def: &StmtFunction
         return;
     }
 
+    if function_def.is_async {
+        return;
+    }
+
     let mut visitor = ReturnInGeneratorVisitor::default();
     visitor.visit_body(&function_def.body);
 

diff --git a/.../ruff_linter__linter__tests__semantic_syntax_error_return_in_async_generator.py_3.10.snap b/.../ruff_linter__linter__tests__semantic_syntax_error_return_in_async_generator.py_3.10.snap
@@ -8,4 +8,37 @@ invalid-syntax: `return` with value in async generator
 2 |     yield 1
 3 |     return 42
   |     ^^^^^^^^^
+4 |
+5 | def gen(): #ok
   |
+
+invalid-syntax: `return` with value in async generator
+  --> resources/test/fixtures/semantic_errors/return_in_async_generator.py:15:5
+   |
+13 | async def gen():
+14 |     yield 1
+15 |     return foo()
+   |     ^^^^^^^^^^^^
+16 |
+17 | async def gen():
+   |
+
+invalid-syntax: `return` with value in async generator
+  --> resources/test/fixtures/semantic_errors/return_in_async_generator.py:19:5
+   |
+17 | async def gen():
+18 |     yield 1
+19 |     return [1, 2, 3]
+   |     ^^^^^^^^^^^^^^^^
+20 |
+21 | async def gen():
+   |
+
+invalid-syntax: `return` with value in async generator
+  --> resources/test/fixtures/semantic_errors/return_in_async_generator.py:24:5
+   |
+22 |     if True:
+23 |         yield 1
+24 |     return 10
+   |     ^^^^^^^^^
+   |
@@ -3,17 +3,18 @@
 //! 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::statement_visitor;
+use ruff_python_ast::statement_visitor::StatementVisitor;
 use ruff_python_ast::{
     self as ast, Expr, ExprContext, IrrefutablePatternKind, Pattern, PythonVersion, Stmt, StmtExpr,
-    StmtImportFrom,
+    StmtFunctionDef, StmtImportFrom,
     comparable::ComparableExpr,
     helpers,
     visitor::{Visitor, walk_expr},
 };
 use ruff_text_size::{Ranged, TextRange, TextSize};
 use rustc_hash::{FxBuildHasher, FxHashSet};
 use std::fmt::Display;
-
 #[derive(Debug, Default)]
 pub struct SemanticSyntaxChecker {
     /// The checker has traversed past the `__future__` import boundary.
@@ -701,7 +702,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::ReturnInAsyncGenerator,
+                                return_range,
+                            );
+                        }
+                    }
+                }
                 self.seen_futures_boundary = true;
             }
             _ => {
@@ -822,15 +837,7 @@ impl SemanticSyntaxChecker {
                     // def f(): yield *x
                     Self::invalid_star_expression(value, ctx);
                 }
-                if ctx.in_function_scope() && ctx.in_async_context() {
-                    if let Some(return_range) = ctx.has_return() {
-                        Self::add_error(
-                            ctx,
-                            SemanticSyntaxErrorKind::ReturnInAsyncGenerator,
-                            return_range,
-                        );
-                    }
-                }
+
                 Self::yield_outside_function(ctx, expr, YieldOutsideFunctionKind::Yield);
             }
             Expr::YieldFrom(_) => {
@@ -1409,10 +1416,7 @@ pub enum SemanticSyntaxErrorKind {
     /// to be very rare and not worth the additional complexity to detect.
     ///
     /// [#111123]: https://github.com/python/cpython/issues/111123
-    LoadBeforeGlobalDeclaration {
-        name: String,
-        start: TextSize,
-    },
+    LoadBeforeGlobalDeclaration { name: String, start: TextSize },
 
     /// Represents the use of a `nonlocal` variable before its `nonlocal` declaration.
     ///
@@ -1430,10 +1434,7 @@ pub enum SemanticSyntaxErrorKind {
     /// ## Known Issues
     ///
     /// See [`LoadBeforeGlobalDeclaration`][Self::LoadBeforeGlobalDeclaration].
-    LoadBeforeNonlocalDeclaration {
-        name: String,
-        start: TextSize,
-    },
+    LoadBeforeNonlocalDeclaration { name: String, start: TextSize },
 
     /// Represents the use of a starred expression in an invalid location, such as a `return` or
     /// `yield` statement.
@@ -1591,6 +1592,7 @@ pub enum SemanticSyntaxErrorKind {
     /// Represents a nonlocal statement for a name that has no binding in an enclosing scope.
     NonlocalWithoutBinding(String),
 
+    /// Represents a `return` statement with a value in an asynchronous generator.
     ReturnInAsyncGenerator,
 }
 
@@ -1708,6 +1710,36 @@ impl Visitor<'_> for ReboundComprehensionVisitor<'_> {
     }
 }
 
+#[derive(Default)]
+struct ReturnVisitor {
+    return_range: Option<TextRange>,
+    has_yield: bool,
+}
+
+impl StatementVisitor<'_> for ReturnVisitor {
+    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(_) | Stmt::ClassDef(_) => {}
+            Stmt::Return(ast::StmtReturn {
+                value: Some(_),
+                range,
+                ..
+            }) => {
+                if self.return_range.is_none() {
+                    self.return_range = Some(*range);
+                }
+            }
+            _ => statement_visitor::walk_stmt(self, stmt),
+        }
+    }
+}
+
 struct MatchPatternVisitor<'a, Ctx> {
     names: FxHashSet<&'a ast::name::Name>,
     ctx: &'a Ctx,
@@ -2128,8 +2160,6 @@ pub trait SemanticSyntaxContext {
     fn in_loop_context(&self) -> bool;
 
     fn is_bound_parameter(&self, name: &str) -> bool;
-
-    fn has_return(&self) -> Option<TextRange>;
 }
 
 /// Modified version of [`std::str::EscapeDefault`] that does not escape single or double quotes.

diff --git a/crates/ruff_python_parser/tests/fixtures.rs b/crates/ruff_python_parser/tests/fixtures.rs
@@ -583,10 +583,6 @@ impl SemanticSyntaxContext for SemanticSyntaxCheckerVisitor<'_> {
     fn is_bound_parameter(&self, _name: &str) -> bool {
         false
     }
-
-    fn has_return(&self) -> Option<TextRange> {
-        None
-    }
 }
 
 impl Visitor<'_> for SemanticSyntaxCheckerVisitor<'_> {

diff --git a/crates/ty_python_semantic/src/semantic_index/builder.rs b/crates/ty_python_semantic/src/semantic_index/builder.rs
@@ -2813,10 +2813,6 @@ impl SemanticSyntaxContext for SemanticIndexBuilder<'_, '_> {
     fn is_bound_parameter(&self, _name: &str) -> bool {
         false
     }
-
-    fn has_return(&self) -> Option<TextRange> {
-        None
-    }
 }
 
 #[derive(Copy, Clone, Debug, PartialEq)]
-Original file line number
+Diff line change
@@ Expand Up @@
             return;
         }
+        if function_def.is_async {
+            return;
+        }
         let mut visitor = ReturnInGeneratorVisitor::default();
         visitor.visit_body(&function_def.body);
@@ Expand Down @@