diff --git a/crates/ruff/resources/test/fixtures/flake8_return/RET504.py b/crates/ruff/resources/test/fixtures/flake8_return/RET504.py index 80b92b3193cf3..9ee07c279cb59 100644 --- a/crates/ruff/resources/test/fixtures/flake8_return/RET504.py +++ b/crates/ruff/resources/test/fixtures/flake8_return/RET504.py @@ -276,20 +276,25 @@ def str_to_bool(val): # Mixed assignments def function_assignment(x): - def f(): ... + def f(): + ... return f def class_assignment(x): - class Foo: ... + class Foo: + ... return Foo def mixed_function_assignment(x): if x: - def f(): ... + + def f(): + ... + else: f = 42 @@ -298,8 +303,32 @@ def f(): ... def mixed_class_assignment(x): if x: - class Foo: ... + + class Foo: + ... + else: Foo = 42 return Foo + + +# `with` statements +def foo(): + with open("foo.txt", "r") as f: + x = f.read() + return x + + +def foo(): + with open("foo.txt", "r") as f: + x = f.read() + print(x) + return x + + +def foo(): + with open("foo.txt", "r") as f: + x = f.read() + print(x) + return x diff --git a/crates/ruff/src/rules/flake8_return/rules/function.rs b/crates/ruff/src/rules/flake8_return/rules/function.rs index c2affd8ed2566..cbf178ded6f02 100644 --- a/crates/ruff/src/rules/flake8_return/rules/function.rs +++ b/crates/ruff/src/rules/flake8_return/rules/function.rs @@ -504,7 +504,7 @@ fn implicit_return(checker: &mut Checker, stmt: &Stmt) { /// RET504 fn unnecessary_assign(checker: &mut Checker, stack: &Stack) { - for (stmt_assign, stmt_return) in &stack.assignments { + for (stmt_assign, stmt_return) in &stack.assignment_return { // Identify, e.g., `return x`. let Some(value) = stmt_return.value.as_ref() else { continue; diff --git a/crates/ruff/src/rules/flake8_return/snapshots/ruff__rules__flake8_return__tests__RET504_RET504.py.snap b/crates/ruff/src/rules/flake8_return/snapshots/ruff__rules__flake8_return__tests__RET504_RET504.py.snap index caa793afe735d..ddf3000678cb8 100644 --- a/crates/ruff/src/rules/flake8_return/snapshots/ruff__rules__flake8_return__tests__RET504_RET504.py.snap +++ b/crates/ruff/src/rules/flake8_return/snapshots/ruff__rules__flake8_return__tests__RET504_RET504.py.snap @@ -41,4 +41,12 @@ RET504.py:268:12: RET504 Unnecessary assignment to `val` before `return` stateme | ^^^ RET504 | +RET504.py:320:12: RET504 Unnecessary assignment to `x` before `return` statement + | +318 | with open("foo.txt", "r") as f: +319 | x = f.read() +320 | return x + | ^ RET504 + | + diff --git a/crates/ruff/src/rules/flake8_return/visitor.rs b/crates/ruff/src/rules/flake8_return/visitor.rs index 4e4bde9ab7b0a..7cabcfe48ca74 100644 --- a/crates/ruff/src/rules/flake8_return/visitor.rs +++ b/crates/ruff/src/rules/flake8_return/visitor.rs @@ -17,7 +17,7 @@ pub(crate) struct Stack<'a> { /// Whether the current function is a generator. pub(crate) is_generator: bool, /// The `assignment`-to-`return` statement pairs in the current function. - pub(crate) assignments: Vec<(&'a ast::StmtAssign, &'a ast::StmtReturn)>, + pub(crate) assignment_return: Vec<(&'a ast::StmtAssign, &'a ast::StmtReturn)>, } #[derive(Default)] @@ -81,8 +81,36 @@ impl<'a> Visitor<'a> for ReturnVisitor<'a> { Stmt::Return(stmt_return) => { // If the `return` statement is preceded by an `assignment` statement, then the // `assignment` statement may be redundant. - if let Some(stmt_assign) = self.sibling.and_then(Stmt::as_assign_stmt) { - self.stack.assignments.push((stmt_assign, stmt_return)); + if let Some(sibling) = self.sibling { + match sibling { + // Example: + // ```python + // def foo(): + // x = 1 + // return x + // ``` + Stmt::Assign(stmt_assign) => { + self.stack + .assignment_return + .push((stmt_assign, stmt_return)); + } + // Example: + // ```python + // def foo(): + // with open("foo.txt", "r") as f: + // x = f.read() + // return x + // ``` + Stmt::With(ast::StmtWith { body, .. }) + | Stmt::AsyncWith(ast::StmtAsyncWith { body, .. }) => { + if let Some(stmt_assign) = body.last().and_then(Stmt::as_assign_stmt) { + self.stack + .assignment_return + .push((stmt_assign, stmt_return)); + } + } + _ => {} + } } self.stack.returns.push(stmt_return);