Add space after return when inlining number for RET504 (#7116)

astral-sh · Sep 3, 2023 · c004e03 · c004e03
1 parent b57ddd5
commit c004e03
Show file tree

Hide file tree

Showing 3 changed files with 30 additions and 3 deletions.
diff --git a/crates/ruff/resources/test/fixtures/flake8_return/RET504.py b/crates/ruff/resources/test/fixtures/flake8_return/RET504.py
@@ -357,3 +357,9 @@ def foo():
 def foo():
     a = 1  # Comment
     return a
+
+
+# Regression test for: https://github.com/astral-sh/ruff/issues/7098
+def mavko_debari(P_kbar):
+    D=0.4853881 + 3.6006116*P - 0.0117368*(P-1.3822)**2
+    return D
diff --git a/crates/ruff/src/rules/flake8_return/rules/function.rs b/crates/ruff/src/rules/flake8_return/rules/function.rs
@@ -12,6 +12,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_semantic::SemanticModel;
+use ruff_python_trivia::is_python_whitespace;
 
 use crate::autofix::edits;
 use crate::checkers::ast::Checker;
@@ -549,11 +550,11 @@ fn unnecessary_assign(checker: &mut Checker, stack: &Stack) {
                     if content[after_equals..]
                         .chars()
                         .next()
-                        .is_some_and(char::is_alphabetic)
+                        .is_some_and(is_python_whitespace)
                     {
-                        "return ".to_string()
-                    } else {
                         "return".to_string()
+                    } else {
+                        "return ".to_string()
                     },
                     // Replace from the start of the assignment statement to the end of the equals
                     // sign.

diff --git a/...rc/rules/flake8_return/snapshots/ruff__rules__flake8_return__tests__RET504_RET504.py.snap b/...rc/rules/flake8_return/snapshots/ruff__rules__flake8_return__tests__RET504_RET504.py.snap
@@ -197,5 +197,25 @@ RET504.py:359:12: RET504 [*] Unnecessary assignment to `a` before `return` state
 358     |-    a = 1  # Comment
 359     |-    return a
     358 |+    return 1  # Comment
+360 359 | 
+361 360 | 
+362 361 | # Regression test for: https://github.com/astral-sh/ruff/issues/7098
+
+RET504.py:365:12: RET504 [*] Unnecessary assignment to `D` before `return` statement
+    |
+363 | def mavko_debari(P_kbar):
+364 |     D=0.4853881 + 3.6006116*P - 0.0117368*(P-1.3822)**2
+365 |     return D
+    |            ^ RET504
+    |
+    = help: Remove unnecessary assignment
+
+ℹ Suggested fix
+361 361 | 
+362 362 | # Regression test for: https://github.com/astral-sh/ruff/issues/7098
+363 363 | def mavko_debari(P_kbar):
+364     |-    D=0.4853881 + 3.6006116*P - 0.0117368*(P-1.3822)**2
+365     |-    return D
+    364 |+    return 0.4853881 + 3.6006116*P - 0.0117368*(P-1.3822)**2