Respect mixed variable assignment in RET504 (#4835)

crates/ruff/resources/test/fixtures/flake8_return/RET504.py

@@ -272,3 +272,34 @@ def str_to_bool(val):
     if isinstance(val, bool):
         return some_obj
     return val
+
+
+# Mixed assignments
+def function_assignment(x):
+    def f(): ...
+
+    return f
+
+
+def class_assignment(x):
+    class Foo: ...
+
+    return Foo
+
+
+def mixed_function_assignment(x):
+    if x:
+        def f(): ...
+    else:
+        f = 42
+
+    return f
+
+
+def mixed_class_assignment(x):
+    if x:
+        class Foo: ...
+    else:
+        Foo = 42
+
+    return Foo

crates/ruff/src/rules/flake8_return/rules/function.rs

@@ -506,34 +506,36 @@ fn implicit_return(checker: &mut Checker, stmt: &Stmt) {
     }
 }
 
-/// Return `true` if the `id` has multiple assignments within the function.
-fn has_multiple_assigns(id: &str, stack: &Stack) -> bool {
-    if let Some(assigns) = stack.assignments.get(&id) {
-        if assigns.len() > 1 {
-            return true;
-        }
-    }
-    false
+/// Return `true` if the `id` has multiple declarations within the function.
+fn has_multiple_declarations(id: &str, stack: &Stack) -> bool {
+    stack
+        .declarations
+        .get(&id)
+        .map_or(false, |declarations| declarations.len() > 1)
 }
 
 /// Return `true` if the `id` has a (read) reference between the `return_location` and its
-/// preceding assignment.
-fn has_refs_before_next_assign(id: &str, return_range: TextRange, stack: &Stack) -> bool {
-    let mut assignment_before_return: Option<TextSize> = None;
-    let mut assignment_after_return: Option<TextSize> = None;
-    if let Some(assignments) = stack.assignments.get(&id) {
+/// preceding declaration.
+fn has_references_before_next_declaration(
+    id: &str,
+    return_range: TextRange,
+    stack: &Stack,
+) -> bool {
+    let mut declaration_before_return: Option<TextSize> = None;
+    let mut declaration_after_return: Option<TextSize> = None;
+    if let Some(assignments) = stack.declarations.get(&id) {
         for location in assignments.iter().sorted() {
             if *location > return_range.start() {
-                assignment_after_return = Some(*location);
+                declaration_after_return = Some(*location);
                 break;
             }
-            assignment_before_return = Some(*location);
+            declaration_before_return = Some(*location);
         }
     }
 
-    // If there is no assignment before the return, then the variable must be defined in
+    // If there is no declaration before the return, then the variable must be declared in
     // some other way (e.g., a function argument). No need to check for references.
-    let Some(assignment_before_return) = assignment_before_return else {
+    let Some(declaration_before_return) = declaration_before_return else {
         return true;
     };
 
@@ -543,9 +545,9 @@ fn has_refs_before_next_assign(id: &str, return_range: TextRange, stack: &Stack)
                 continue;
             }
 
-            if assignment_before_return < *location {
-                if let Some(assignment_after_return) = assignment_after_return {
-                    if *location <= assignment_after_return {
+            if declaration_before_return < *location {
+                if let Some(declaration_after_return) = declaration_after_return {
+                    if *location <= declaration_after_return {
                         return true;
                     }
                 } else {
@@ -559,7 +561,7 @@ fn has_refs_before_next_assign(id: &str, return_range: TextRange, stack: &Stack)
 }
 
 /// Return `true` if the `id` has a read or write reference within a `try` or loop body.
-fn has_refs_or_assigns_within_try_or_loop(id: &str, stack: &Stack) -> bool {
+fn has_references_or_declarations_within_try_or_loop(id: &str, stack: &Stack) -> bool {
     if let Some(references) = stack.references.get(&id) {
         for location in references {
             for try_range in &stack.tries {
@@ -574,7 +576,7 @@ fn has_refs_or_assigns_within_try_or_loop(id: &str, stack: &Stack) -> bool {
             }
         }
     }
-    if let Some(references) = stack.assignments.get(&id) {
+    if let Some(references) = stack.declarations.get(&id) {
         for location in references {
             for try_range in &stack.tries {
                 if try_range.contains(*location) {
@@ -594,7 +596,7 @@ fn has_refs_or_assigns_within_try_or_loop(id: &str, stack: &Stack) -> bool {
 /// RET504
 fn unnecessary_assign(checker: &mut Checker, stack: &Stack, expr: &Expr) {
     if let Expr::Name(ast::ExprName { id, .. }) = expr {
-        if !stack.assignments.contains_key(id.as_str()) {
+        if !stack.assigned_names.contains(id.as_str()) {
             return;
         }
 
@@ -605,9 +607,9 @@ fn unnecessary_assign(checker: &mut Checker, stack: &Stack, expr: &Expr) {
             return;
         }
 
-        if has_multiple_assigns(id, stack)
-            || has_refs_before_next_assign(id, expr.range(), stack)
-            || has_refs_or_assigns_within_try_or_loop(id, stack)
+        if has_multiple_declarations(id, stack)
+            || has_references_before_next_declaration(id, expr.range(), stack)
+            || has_references_or_declarations_within_try_or_loop(id, stack)
         {
             return;
         }

crates/ruff/src/rules/flake8_return/visitor.rs

@@ -11,9 +11,14 @@ pub(crate) struct Stack<'a> {
     pub(crate) yields: Vec<&'a Expr>,
     pub(crate) elses: Vec<&'a Stmt>,
     pub(crate) elifs: Vec<&'a Stmt>,
+    /// The names that are assigned to in the current scope (e.g., anything on the left-hand side of
+    /// an assignment).
+    pub(crate) assigned_names: FxHashSet<&'a str>,
+    /// The names that are declared in the current scope, and the ranges of those declarations
+    /// (e.g., assignments, but also function and class definitions).
+    pub(crate) declarations: FxHashMap<&'a str, Vec<TextSize>>,
     pub(crate) references: FxHashMap<&'a str, Vec<TextSize>>,
     pub(crate) non_locals: FxHashSet<&'a str>,
-    pub(crate) assignments: FxHashMap<&'a str, Vec<TextSize>>,
     pub(crate) loops: Vec<TextRange>,
     pub(crate) tries: Vec<TextRange>,
 }
@@ -34,8 +39,9 @@ impl<'a> ReturnVisitor<'a> {
                 return;
             }
             Expr::Name(ast::ExprName { id, .. }) => {
+                self.stack.assigned_names.insert(id.as_str());
                 self.stack
-                    .assignments
+                    .declarations
                     .entry(id)
                     .or_insert_with(Vec::new)
                     .push(expr.start());
@@ -45,7 +51,7 @@ impl<'a> ReturnVisitor<'a> {
                 // Attribute assignments are often side-effects (e.g., `self.property = value`),
                 // so we conservatively treat them as references to every known
                 // variable.
-                for name in self.stack.assignments.keys() {
+                for name in self.stack.declarations.keys() {
                     self.stack
                         .references
                         .entry(name)
@@ -68,18 +74,44 @@ impl<'a> Visitor<'a> for ReturnVisitor<'a> {
                     .non_locals
                     .extend(names.iter().map(Identifier::as_str));
             }
-            Stmt::FunctionDef(ast::StmtFunctionDef {
+            Stmt::ClassDef(ast::StmtClassDef {
                 decorator_list,
+                name,
+                ..
+            }) => {
+                // Mark a declaration.
+                self.stack
+                    .declarations
+                    .entry(name.as_str())
+                    .or_insert_with(Vec::new)
+                    .push(stmt.start());
+
+                // Don't recurse into the body, but visit the decorators, etc.
+                for expr in decorator_list {
+                    visitor::walk_expr(self, expr);
+                }
+            }
+            Stmt::FunctionDef(ast::StmtFunctionDef {
+                name,
                 args,
+                decorator_list,
                 returns,
                 ..
             })
             | Stmt::AsyncFunctionDef(ast::StmtAsyncFunctionDef {
-                decorator_list,
+                name,
                 args,
+                decorator_list,
                 returns,
                 ..
             }) => {
+                // Mark a declaration.
+                self.stack
+                    .declarations
+                    .entry(name.as_str())
+                    .or_insert_with(Vec::new)
+                    .push(stmt.start());
+
                 // Don't recurse into the body, but visit the decorators, etc.
                 for expr in decorator_list {
                     visitor::walk_expr(self, expr);
@@ -138,7 +170,7 @@ impl<'a> Visitor<'a> for ReturnVisitor<'a> {
 
                 if let Some(target) = targets.first() {
                     // Skip unpacking assignments, like `x, y = my_object`.
-                    if matches!(target, Expr::Tuple(_)) && !value.is_tuple_expr() {
+                    if target.is_tuple_expr() && !value.is_tuple_expr() {
                         return;
                     }
 
@@ -172,7 +204,7 @@ impl<'a> Visitor<'a> for ReturnVisitor<'a> {
             Expr::Call(_) => {
                 // Arbitrary function calls can have side effects, so we conservatively treat
                 // every function call as a reference to every known variable.
-                for name in self.stack.assignments.keys() {
+                for name in self.stack.declarations.keys() {
                     self.stack
                         .references
                         .entry(name)