diff --git a/crates/ruff/resources/test/fixtures/flake8_return/RET504.py b/crates/ruff/resources/test/fixtures/flake8_return/RET504.py index 96bc755687..80b92b3193 100644 --- a/crates/ruff/resources/test/fixtures/flake8_return/RET504.py +++ b/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 diff --git a/crates/ruff/src/rules/flake8_return/rules/function.rs b/crates/ruff/src/rules/flake8_return/rules/function.rs index cc84657be0..9bfb85ade3 100644 --- a/crates/ruff/src/rules/flake8_return/rules/function.rs +++ b/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 = None; - let mut assignment_after_return: Option = 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 = None; + let mut declaration_after_return: Option = 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; } diff --git a/crates/ruff/src/rules/flake8_return/visitor.rs b/crates/ruff/src/rules/flake8_return/visitor.rs index c7c3657138..f0d419c7bc 100644 --- a/crates/ruff/src/rules/flake8_return/visitor.rs +++ b/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>, pub(crate) references: FxHashMap<&'a str, Vec>, pub(crate) non_locals: FxHashSet<&'a str>, - pub(crate) assignments: FxHashMap<&'a str, Vec>, pub(crate) loops: Vec, pub(crate) tries: Vec, } @@ -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)