From 42c8054268d16a062919879cd543e7385a01ca03 Mon Sep 17 00:00:00 2001 From: Charlie Marsh Date: Sat, 10 Jun 2023 00:32:03 -0400 Subject: [PATCH] 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. --- .../test/fixtures/flake8_return/RET504.py | 31 ++- .../src/rules/flake8_return/rules/function.rs | 69 ++++++- ...lake8_return__tests__RET504_RET504.py.snap | 185 ++++++++++++++++-- .../ruff/src/rules/flake8_return/visitor.rs | 8 +- 4 files changed, 262 insertions(+), 31 deletions(-) diff --git a/crates/ruff/resources/test/fixtures/flake8_return/RET504.py b/crates/ruff/resources/test/fixtures/flake8_return/RET504.py index 9ee07c279c..7bd789386a 100644 --- a/crates/ruff/resources/test/fixtures/flake8_return/RET504.py +++ b/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 diff --git a/crates/ruff/src/rules/flake8_return/rules/function.rs b/crates/ruff/src/rules/flake8_return/rules/function.rs index cbf178ded6..8f5e3aaf88 100644 --- a/crates/ruff/src/rules/flake8_return/rules/function.rs +++ b/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); } } 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 ddf3000678..ae04abcf8c 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 @@ -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 = queryset.filter(c=3) -245 | return queryset +244 | queryset = Model.filter(a=1) +245 | queryset = queryset.filter(c=3) +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 | queryset = Model.filter(a=1) -250 | return queryset # RET504 +249 | def get_queryset(): +250 | queryset = Model.filter(a=1) +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 | val = 1 -268 | return val # RET504 +267 | return val +268 | val = 1 +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 | x = f.read() -320 | return x +319 | with open("foo.txt", "r") as f: +320 | x = f.read() +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 diff --git a/crates/ruff/src/rules/flake8_return/visitor.rs b/crates/ruff/src/rules/flake8_return/visitor.rs index 7cabcfe48c..d07c12b184 100644 --- a/crates/ruff/src/rules/flake8_return/visitor.rs +++ b/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)); } } _ => {}