[flake8-return] Fix RET504 autofix generating a syntax error (#18428)

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

@@ -421,3 +421,14 @@ def func(a: dict[str, int]) -> list[dict[str, int]]:
     if "services" in a:
         services = a["services"]
         return services
+
+# See: https://github.com/astral-sh/ruff/issues/18411
+def f():
+    (#=
+    x) = 1
+    return x
+
+def f():
+    x = (1
+              )
+    return x

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

@@ -1,5 +1,3 @@
-use std::ops::Add;
-
 use anyhow::Result;
 
 use ruff_macros::{ViolationMetadata, derive_message_formats};
@@ -8,6 +6,7 @@ use ruff_python_ast::stmt_if::elif_else_range;
 use ruff_python_ast::visitor::Visitor;
 use ruff_python_ast::whitespace::indentation;
 use ruff_python_ast::{self as ast, Decorator, ElifElseClause, Expr, Stmt};
+use ruff_python_parser::TokenKind;
 use ruff_python_semantic::SemanticModel;
 use ruff_python_semantic::analyze::visibility::is_property;
 use ruff_python_trivia::{SimpleTokenKind, SimpleTokenizer, is_python_whitespace};
@@ -614,17 +613,17 @@ fn unnecessary_assign(checker: &Checker, stack: &Stack) {
             let delete_return =
                 edits::delete_stmt(stmt, None, checker.locator(), checker.indexer());
 
-            // Replace the `x = 1` statement with `return 1`.
-            let content = checker.locator().slice(assign);
-            let equals_index = content
-                .find('=')
-                .ok_or(anyhow::anyhow!("expected '=' in assignment statement"))?;
-            let after_equals = equals_index + 1;
+            let eq_token = checker
+                .tokens()
+                .before(assign.value.start())
+                .iter()
+                .rfind(|token| token.kind() == TokenKind::Equal)
+                .unwrap();
 
+            let content = checker.source();
+            // Replace the `x = 1` statement with `return 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..]
+                if content[eq_token.end().to_usize()..]
                     .chars()
                     .next()
                     .is_some_and(is_python_whitespace)
@@ -635,13 +634,7 @@ fn unnecessary_assign(checker: &Checker, stack: &Stack) {
                 },
                 // Replace from the start of the assignment statement to the end of the equals
                 // sign.
-                TextRange::new(
-                    assign.start(),
-                    assign
-                        .range()
-                        .start()
-                        .add(TextSize::try_from(after_equals)?),
-                ),
+                TextRange::new(assign.start(), eq_token.range().end()),
             );
 
             Ok(Fix::unsafe_edits(replace_assign, [delete_return]))

crates/ruff_linter/src/rules/flake8_return/snapshots/ruff_linter__rules__flake8_return__tests__RET504_RET504.py.snap

@@ -1,6 +1,5 @@
 ---
 source: crates/ruff_linter/src/rules/flake8_return/mod.rs
-snapshot_kind: text
 ---
 RET504.py:6:12: RET504 [*] Unnecessary assignment to `a` before `return` statement
   |
@@ -248,6 +247,8 @@ RET504.py:423:16: RET504 [*] Unnecessary assignment to `services` before `return
 422 |         services = a["services"]
 423 |         return services
     |                ^^^^^^^^ RET504
+424 |
+425 | # See: https://github.com/astral-sh/ruff/issues/18411
     |
     = help: Remove unnecessary assignment
 
@@ -258,3 +259,47 @@ RET504.py:423:16: RET504 [*] Unnecessary assignment to `services` before `return
 422     |-        services = a["services"]
 423     |-        return services
     422 |+        return a["services"]
+424 423 | 
+425 424 | # See: https://github.com/astral-sh/ruff/issues/18411
+426 425 | def f():
+
+RET504.py:429:12: RET504 [*] Unnecessary assignment to `x` before `return` statement
+    |
+427 |     (#=
+428 |     x) = 1
+429 |     return x
+    |            ^ RET504
+430 |
+431 | def f():
+    |
+    = help: Remove unnecessary assignment
+
+ℹ Unsafe fix
+424 424 | 
+425 425 | # See: https://github.com/astral-sh/ruff/issues/18411
+426 426 | def f():
+427     |-    (#=
+428     |-    x) = 1
+429     |-    return x
+    427 |+    return 1
+430 428 | 
+431 429 | def f():
+432 430 |     x = (1
+
+RET504.py:434:12: RET504 [*] Unnecessary assignment to `x` before `return` statement
+    |
+432 |     x = (1
+433 |               )
+434 |     return x
+    |            ^ RET504
+    |
+    = help: Remove unnecessary assignment
+
+ℹ Unsafe fix
+429 429 |     return x
+430 430 | 
+431 431 | def f():
+432     |-    x = (1
+    432 |+    return (1
+433 433 |               )
+434     |-    return x