Extend revised RET504 implementation to with statements (#4998)

## Summary

This PR extends the new `RET504` implementation to handle cases like:

```py
def foo():
    with open("foo.txt", "r") as f:
        x = f.read()
    return x
```

This was originally suggested in
https://github.com/astral-sh/ruff/issues/2950#issuecomment-1433441503.

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 mixed_function_assignment(x):
 
 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

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;

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
+    |
+
 

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) {
+                if let Some(sibling) = self.sibling {
-                    self.stack.assignments.push((stmt_assign, stmt_return));
+                    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);