From 8f6912a519b44ba266d2b94c9befcd2d595c5769 Mon Sep 17 00:00:00 2001
From: Charlie Marsh <charlie.r.marsh@gmail.com>
Date: Fri, 9 Jun 2023 22:51:43 -0400
Subject: [PATCH] Extend RET504 to with statements

---
 .../test/fixtures/flake8_return/RET504.py     | 37 +++++++++++++++++--
 .../src/rules/flake8_return/rules/function.rs |  2 +-
 ...lake8_return__tests__RET504_RET504.py.snap |  8 ++++
 .../ruff/src/rules/flake8_return/visitor.rs   | 34 +++++++++++++++--
 4 files changed, 73 insertions(+), 8 deletions(-)

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);