Implement autofix for revised `RET504` rule (#4999)

## Summary

This PR enables autofix for the revised `RET504` rule, by changing:

```py
def f():
    x = 1
    return x
```

...to:

```py
def f():
    return 1
```

Closes #2263.

Closes #2788.

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

@@ -79,7 +79,7 @@ def x():
     return a
 
 
-# ignore unpacking
+# Ignore unpacking
 def x():
     b, a = [1, 2]
     return a
@@ -109,7 +109,8 @@ def x():
 
 # Considered OK, since functions can have side effects.
 def x():
-    b, a = 1, 2
+    a = 1
+    b = 2
     print(b)
     return a
 
@@ -317,7 +318,7 @@ def mixed_class_assignment(x):
 def foo():
     with open("foo.txt", "r") as f:
         x = f.read()
-    return x
+    return x  # RET504
 
 
 def foo():
@@ -332,3 +333,27 @@ def foo():
         x = f.read()
     print(x)
     return x
+
+
+# Autofix cases
+def foo():
+    a = 1
+    b=a
+    return b  # RET504
+
+
+def foo():
+    a = 1
+    b =a
+    return b  # RET504
+
+
+def foo():
+    a = 1
+    b= a
+    return b  # RET504
+
+
+def foo():
+    a = 1  # Comment
+    return a

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

@@ -1,3 +1,6 @@
+use std::ops::Add;
+
+use ruff_text_size::{TextRange, TextSize};
 use rustpython_parser::ast::{self, Constant, Expr, Ranged, Stmt};
 
 use ruff_diagnostics::{AlwaysAutofixableViolation, Violation};
@@ -9,6 +12,7 @@ use ruff_python_ast::visitor::Visitor;
 use ruff_python_ast::whitespace::indentation;
 use ruff_python_semantic::model::SemanticModel;
 
+use crate::autofix::edits;
 use crate::checkers::ast::Checker;
 use crate::registry::{AsRule, Rule};
 use crate::rules::flake8_return::helpers::end_of_last_statement;
@@ -161,12 +165,16 @@ pub struct UnnecessaryAssign {
     name: String,
 }
 
-impl Violation for UnnecessaryAssign {
+impl AlwaysAutofixableViolation for UnnecessaryAssign {
     #[derive_message_formats]
     fn message(&self) -> String {
         let UnnecessaryAssign { name } = self;
         format!("Unnecessary assignment to `{name}` before `return` statement")
     }
+
+    fn autofix_title(&self) -> String {
+        "Remove unnecessary assignment".to_string()
+    }
 }
 
 /// ## What it does
@@ -504,9 +512,9 @@ fn implicit_return(checker: &mut Checker, stmt: &Stmt) {
 
 /// RET504
 fn unnecessary_assign(checker: &mut Checker, stack: &Stack) {
-    for (stmt_assign, stmt_return) in &stack.assignment_return {
+    for (assign, return_, stmt) in &stack.assignment_return {
         // Identify, e.g., `return x`.
-        let Some(value) = stmt_return.value.as_ref() else {
+        let Some(value) = return_.value.as_ref() else {
             continue;
         };
 
@@ -515,11 +523,11 @@ fn unnecessary_assign(checker: &mut Checker, stack: &Stack) {
         };
 
         // Identify, e.g., `x = 1`.
-        if stmt_assign.targets.len() > 1 {
+        if assign.targets.len() > 1 {
             continue;
         }
 
-        let Some(target) = stmt_assign.targets.first() else {
+        let Some(target) = assign.targets.first() else {
             continue;
         };
 
@@ -535,12 +543,59 @@ fn unnecessary_assign(checker: &mut Checker, stack: &Stack) {
             continue;
         }
 
-        checker.diagnostics.push(Diagnostic::new(
+        let mut diagnostic = Diagnostic::new(
             UnnecessaryAssign {
                 name: assigned_id.to_string(),
             },
             value.range(),
-        ));
+        );
+        if checker.patch(diagnostic.kind.rule()) {
+            diagnostic.try_set_fix(|| {
+                // Delete the `return` statement. There's no need to treat this as an isolated
+                // edit, since we're editing the preceding statement, so no conflicting edit would
+                // be allowed to remove that preceding statement.
+                let delete_return = edits::delete_stmt(
+                    stmt,
+                    None,
+                    checker.locator,
+                    checker.indexer,
+                    checker.stylist,
+                );
+
+                // Replace the `x = 1` statement with `return 1`.
+                let content = checker.locator.slice(assign.range());
+                let equals_index = content
+                    .find('=')
+                    .ok_or(anyhow::anyhow!("expected '=' in assignment statement"))?;
+                let after_equals = equals_index + 1;
+
+                let replace_assign = Edit::range_replacement(
+                    // If necessary, add whitespace after the `return` keyword.
+                    // Ex) Convert `x=y` to `return y` (instead of `returny`).
+                    if content[after_equals..]
+                        .chars()
+                        .next()
+                        .map_or(false, char::is_alphabetic)
+                    {
+                        "return ".to_string()
+                    } else {
+                        "return".to_string()
+                    },
+                    // Replace from the start of the assignment statement to the end of the equals
+                    // sign.
+                    TextRange::new(
+                        assign.range().start(),
+                        assign
+                            .range()
+                            .start()
+                            .add(TextSize::try_from(after_equals)?),
+                    ),
+                );
+
+                Ok(Fix::suggested_edits(replace_assign, [delete_return]))
+            });
+        }
+        checker.diagnostics.push(diagnostic);
     }
 }
 

crates/ruff/src/rules/flake8_return/snapshots/ruff__rules__flake8_return__tests__RET504_RET504.py.snap

@@ -1,52 +1,201 @@
 ---
 source: crates/ruff/src/rules/flake8_return/mod.rs
 ---
-RET504.py:6:12: RET504 Unnecessary assignment to `a` before `return` statement
+RET504.py:6:12: RET504 [*] Unnecessary assignment to `a` before `return` statement
   |
 4 | def x():
 5 |     a = 1
 6 |     return a  # RET504
   |            ^ RET504
   |
+  = help: Remove unnecessary assignment
 
-RET504.py:23:12: RET504 Unnecessary assignment to `formatted` before `return` statement
+ℹ Suggested fix
+2 2 | # Errors
+3 3 | ###
+4 4 | def x():
+5   |-    a = 1
+6   |-    return a  # RET504
+  5 |+    return 1
+7 6 | 
+8 7 | 
+9 8 | # Can be refactored false positives
+
+RET504.py:23:12: RET504 [*] Unnecessary assignment to `formatted` before `return` statement
    |
 21 |     # clean up after any blank components
 22 |     formatted = formatted.replace("()", "").replace("  ", " ").strip()
 23 |     return formatted
    |            ^^^^^^^^^ RET504
    |
+   = help: Remove unnecessary assignment
 
-RET504.py:245:12: RET504 Unnecessary assignment to `queryset` before `return` statement
+ℹ Suggested fix
+19 19 | def x():
+20 20 |     formatted = _USER_AGENT_FORMATTER.format(format_string, **values)
+21 21 |     # clean up after any blank components
+22    |-    formatted = formatted.replace("()", "").replace("  ", " ").strip()
+23    |-    return formatted
+   22 |+    return formatted.replace("()", "").replace("  ", " ").strip()
+24 23 | 
+25 24 | 
+26 25 | # https://github.com/afonasev/flake8-return/issues/47#issue-641117366
+
+RET504.py:246:12: RET504 [*] Unnecessary assignment to `queryset` before `return` statement
     |
-243 |     queryset = Model.filter(a=1)
+244 |     queryset = Model.filter(a=1)
-244 |     queryset = queryset.filter(c=3)
+245 |     queryset = queryset.filter(c=3)
-245 |     return queryset
+246 |     return queryset
     |            ^^^^^^^^ RET504
     |
+    = help: Remove unnecessary assignment
 
-RET504.py:250:12: RET504 Unnecessary assignment to `queryset` before `return` statement
+ℹ Suggested fix
+242 242 | 
+243 243 | def get_queryset():
+244 244 |     queryset = Model.filter(a=1)
+245     |-    queryset = queryset.filter(c=3)
+246     |-    return queryset
+    245 |+    return queryset.filter(c=3)
+247 246 | 
+248 247 | 
+249 248 | def get_queryset():
+
+RET504.py:251:12: RET504 [*] Unnecessary assignment to `queryset` before `return` statement
     |
-248 | def get_queryset():
+249 | def get_queryset():
-249 |     queryset = Model.filter(a=1)
+250 |     queryset = Model.filter(a=1)
-250 |     return queryset  # RET504
+251 |     return queryset  # RET504
     |            ^^^^^^^^ RET504
     |
+    = help: Remove unnecessary assignment
 
-RET504.py:268:12: RET504 Unnecessary assignment to `val` before `return` statement
+ℹ Suggested fix
+247 247 | 
+248 248 | 
+249 249 | def get_queryset():
+250     |-    queryset = Model.filter(a=1)
+251     |-    return queryset  # RET504
+    250 |+    return Model.filter(a=1)
+252 251 | 
+253 252 | 
+254 253 | # Function arguments
+
+RET504.py:269:12: RET504 [*] Unnecessary assignment to `val` before `return` statement
     |
-266 |         return val
+267 |         return val
-267 |     val = 1
+268 |     val = 1
-268 |     return val  # RET504
+269 |     return val  # RET504
     |            ^^^ RET504
     |
+    = help: Remove unnecessary assignment
 
-RET504.py:320:12: RET504 Unnecessary assignment to `x` before `return` statement
+ℹ Suggested fix
+265 265 | def str_to_bool(val):
+266 266 |     if isinstance(val, bool):
+267 267 |         return val
+268     |-    val = 1
+269     |-    return val  # RET504
+    268 |+    return 1
+270 269 | 
+271 270 | 
+272 271 | def str_to_bool(val):
+
+RET504.py:321:12: RET504 [*] Unnecessary assignment to `x` before `return` statement
     |
-318 |     with open("foo.txt", "r") as f:
+319 |     with open("foo.txt", "r") as f:
-319 |         x = f.read()
+320 |         x = f.read()
-320 |     return x
+321 |     return x  # RET504
     |            ^ RET504
     |
+    = help: Remove unnecessary assignment
+
+ℹ Suggested fix
+317 317 | # `with` statements
+318 318 | def foo():
+319 319 |     with open("foo.txt", "r") as f:
+320     |-        x = f.read()
+321     |-    return x  # RET504
+    320 |+        return f.read()
+322 321 | 
+323 322 | 
+324 323 | def foo():
+
+RET504.py:342:12: RET504 [*] Unnecessary assignment to `b` before `return` statement
+    |
+340 |     a = 1
+341 |     b=a
+342 |     return b  # RET504
+    |            ^ RET504
+    |
+    = help: Remove unnecessary assignment
+
+ℹ Suggested fix
+338 338 | # Autofix cases
+339 339 | def foo():
+340 340 |     a = 1
+341     |-    b=a
+342     |-    return b  # RET504
+    341 |+    return a
+343 342 | 
+344 343 | 
+345 344 | def foo():
+
+RET504.py:348:12: RET504 [*] Unnecessary assignment to `b` before `return` statement
+    |
+346 |     a = 1
+347 |     b =a
+348 |     return b  # RET504
+    |            ^ RET504
+    |
+    = help: Remove unnecessary assignment
+
+ℹ Suggested fix
+344 344 | 
+345 345 | def foo():
+346 346 |     a = 1
+347     |-    b =a
+348     |-    return b  # RET504
+    347 |+    return a
+349 348 | 
+350 349 | 
+351 350 | def foo():
+
+RET504.py:354:12: RET504 [*] Unnecessary assignment to `b` before `return` statement
+    |
+352 |     a = 1
+353 |     b= a
+354 |     return b  # RET504
+    |            ^ RET504
+    |
+    = help: Remove unnecessary assignment
+
+ℹ Suggested fix
+350 350 | 
+351 351 | def foo():
+352 352 |     a = 1
+353     |-    b= a
+354     |-    return b  # RET504
+    353 |+    return a
+355 354 | 
+356 355 | 
+357 356 | def foo():
+
+RET504.py:359:12: RET504 [*] Unnecessary assignment to `a` before `return` statement
+    |
+357 | def foo():
+358 |     a = 1  # Comment
+359 |     return a
+    |            ^ RET504
+    |
+    = help: Remove unnecessary assignment
+
+ℹ Suggested fix
+355 355 | 
+356 356 | 
+357 357 | def foo():
+358     |-    a = 1  # Comment
+359     |-    return a
+    358 |+    return 1  # Comment
 
 

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

@@ -17,7 +17,9 @@ 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) assignment_return: Vec<(&'a ast::StmtAssign, &'a ast::StmtReturn)>,
+    /// TODO(charlie): Remove the extra [`Stmt`] here, which is necessary to support statement
+    /// removal for the `return` statement.
+    pub(crate) assignment_return: Vec<(&'a ast::StmtAssign, &'a ast::StmtReturn, &'a Stmt)>,
 }
 
 #[derive(Default)]
@@ -92,7 +94,7 @@ impl<'a> Visitor<'a> for ReturnVisitor<'a> {
                         Stmt::Assign(stmt_assign) => {
                             self.stack
                                 .assignment_return
-                                .push((stmt_assign, stmt_return));
+                                .push((stmt_assign, stmt_return, stmt));
                         }
                         // Example:
                         // ```python
@@ -106,7 +108,7 @@ impl<'a> Visitor<'a> for ReturnVisitor<'a> {
                             if let Some(stmt_assign) = body.last().and_then(Stmt::as_assign_stmt) {
                                 self.stack
                                     .assignment_return
-                                    .push((stmt_assign, stmt_return));
+                                    .push((stmt_assign, stmt_return, stmt));
                             }
                         }
                         _ => {}