diff --git a/README.md b/README.md index 683dae80c194f..f4e3c360ba707 100644 --- a/README.md +++ b/README.md @@ -1276,7 +1276,7 @@ For more, see [flake8-return](https://pypi.org/project/flake8-return/) on PyPI. | RET501 | unnecessary-return-none | Do not explicitly `return None` in function if it is the only possible return value | 🛠 | | RET502 | implicit-return-value | Do not implicitly `return None` in function able to return non-`None` value | 🛠 | | RET503 | implicit-return | Missing explicit `return` at the end of function able to return non-`None` value | 🛠 | -| RET504 | unnecessary-assign | Unnecessary variable assignment before `return` statement | | +| RET504 | unnecessary-assign | Unnecessary variable assignment before `return` statement | 🛠 | | RET505 | superfluous-else-return | Unnecessary `{branch}` after `return` statement | | | RET506 | superfluous-else-raise | Unnecessary `{branch}` after `raise` statement | | | RET507 | superfluous-else-continue | Unnecessary `{branch}` after `continue` statement | | diff --git a/crates/ruff/resources/test/fixtures/flake8_return/RET504.py b/crates/ruff/resources/test/fixtures/flake8_return/RET504.py index 92bd411d8d0ab..5ccfb302dda93 100644 --- a/crates/ruff/resources/test/fixtures/flake8_return/RET504.py +++ b/crates/ruff/resources/test/fixtures/flake8_return/RET504.py @@ -22,6 +22,18 @@ def x(): return a +def x(): + var = [5, 4, 3, 2, 1] + """ + This is a + block comment + we want to keep + """ + # This is a comment that we want to keep. + return var # Another critical comment + # Last important comment + + # Can be refactored false positives # https://github.com/afonasev/flake8-return/issues/47#issuecomment-1122571066 def get_bar_if_exists(obj): diff --git a/crates/ruff/src/rules/flake8_return/helpers.rs b/crates/ruff/src/rules/flake8_return/helpers.rs index 23e60c200210c..b7d926f3c9f05 100644 --- a/crates/ruff/src/rules/flake8_return/helpers.rs +++ b/crates/ruff/src/rules/flake8_return/helpers.rs @@ -1,5 +1,10 @@ use rustpython_parser::ast::{Constant, Expr, ExprKind, Stmt}; +use rustpython_parser::lexer; +use rustpython_parser::lexer::Tok; + +use crate::source_code::Stylist; + /// Return `true` if a function's return statement include at least one /// non-`None` value. pub fn result_exists(returns: &[(&Stmt, Option<&Expr>)]) -> bool { @@ -16,3 +21,29 @@ pub fn result_exists(returns: &[(&Stmt, Option<&Expr>)]) -> bool { .unwrap_or(false) }) } + +/// Check if the given piece code is composed of exclusively comments +pub fn code_is_only_comments(code: &str) -> bool { + let dedented = textwrap::dedent(code); + let code_tokens = lexer::make_tokenizer(&dedented).flatten(); + for (_, tok, _) in code_tokens { + if !matches!( + tok, + Tok::Comment(..) + | Tok::String { .. } + | Tok::Pass + | Tok::Newline + | Tok::NonLogicalNewline + ) { + return false; + } + } + true +} + +/// Extract the indentation from a given line +pub fn extract_indentation(line: &str, stylist: &Stylist) -> String { + line.chars() + .take_while(|c| stylist.indentation().contains(*c)) + .collect() +} diff --git a/crates/ruff/src/rules/flake8_return/rules.rs b/crates/ruff/src/rules/flake8_return/rules.rs index 6801be192d7e1..33f51ddce0e91 100644 --- a/crates/ruff/src/rules/flake8_return/rules.rs +++ b/crates/ruff/src/rules/flake8_return/rules.rs @@ -1,18 +1,20 @@ use itertools::Itertools; -use ruff_macros::{define_violation, derive_message_formats}; use rustpython_parser::ast::{Constant, Expr, ExprKind, Location, Stmt, StmtKind}; -use super::branch::Branch; -use super::helpers::result_exists; -use super::visitor::{ReturnVisitor, Stack}; -use crate::ast::helpers::elif_else_range; +use ruff_macros::{define_violation, derive_message_formats}; + +use crate::ast::helpers::{elif_else_range, unparse_expr}; use crate::ast::types::Range; use crate::ast::visitor::Visitor; use crate::ast::whitespace::indentation; use crate::checkers::ast::Checker; use crate::fix::Fix; use crate::registry::{Diagnostic, Rule}; -use crate::violation::{AlwaysAutofixableViolation, Violation}; +use crate::violation::{AlwaysAutofixableViolation, AutofixKind, Availability, Violation}; + +use super::branch::Branch; +use super::helpers::{code_is_only_comments, extract_indentation, result_exists}; +use super::visitor::{ReturnVisitor, Stack}; define_violation!( pub struct UnnecessaryReturnNone; @@ -62,10 +64,20 @@ define_violation!( pub struct UnnecessaryAssign; ); impl Violation for UnnecessaryAssign { + const AUTOFIX: Option = Some(AutofixKind { + available: Availability::Sometimes, + }); #[derive_message_formats] fn message(&self) -> String { format!("Unnecessary variable assignment before `return` statement") } + + fn autofix_title_formatter(&self) -> Option String> { + Some(|_| { + "Remove variable definition and place the constant in the `return` statement" + .to_string() + }) + } } define_violation!( @@ -267,7 +279,8 @@ fn has_refs_before_next_assign(id: &str, return_location: Location, stack: &Stac let mut before_assign: &Location = &Location::default(); let mut after_assign: Option<&Location> = None; if let Some(assigns) = stack.assigns.get(&id) { - for location in assigns.iter().sorted() { + for range in assigns.iter().sorted() { + let location = &range.location; if location.row() > return_location.row() { after_assign = Some(location); break; @@ -312,7 +325,8 @@ fn has_refs_or_assigns_within_try_or_loop(id: &str, stack: &Stack) -> bool { } } if let Some(refs) = stack.assigns.get(&id) { - for location in refs { + for range in refs { + let location = &range.location; for (try_location, try_end_location) in &stack.tries { if try_location.row() < location.row() && location.row() <= try_end_location.row() { return true; @@ -337,6 +351,7 @@ fn unnecessary_assign(checker: &mut Checker, stack: &Stack, expr: &Expr) { } if !stack.refs.contains_key(id.as_str()) { + // No autofix, these cases are complicated enough to require human refactoring. checker.diagnostics.push(Diagnostic::new( UnnecessaryAssign, Range::from_located(expr), @@ -354,10 +369,74 @@ fn unnecessary_assign(checker: &mut Checker, stack: &Stack, expr: &Expr) { return; } - checker.diagnostics.push(Diagnostic::new( - UnnecessaryAssign, - Range::from_located(expr), - )); + let mut diagnostic = Diagnostic::new(UnnecessaryAssign, Range::from_located(expr)); + + if checker.patch(diagnostic.kind.rule()) { + // assign_expr is the Expr referring to the value of the last assignment (e.g. `5+5` for `var = 5+5`) + // TODO: This might cause a problem if there is an assignment after, needs further testing + if let Some(assign_expr) = stack.assign_values.get(id.as_str()) { + // assign_locations is a &Vec that has each range for each whole assignment (`var = 5+5`) + let assign_locations = stack.assigns.get(id.as_str()).unwrap(); + let expr_loc = expr.end_location.unwrap(); + if let Some(last_assign_loc) = assign_locations + .iter() + .rev() + .find(|x| x.location < expr_loc) + { + let in_between_source = checker + .locator + .slice(&Range::new(last_assign_loc.end_location, expr.location)) + .to_owned(); + // Ensure that all indentation is identical (no if or while statement). + let all_indentation_same = (checker + .locator + .slice(&Range::new( + Location::new(last_assign_loc.location.row(), 0), + Location::new(last_assign_loc.location.row() + 1, 0), + )) + .chars() + .take_while(|c| checker.stylist.indentation().contains(*c)) + .join("") + + &in_between_source) + .lines() + .tuple_windows() + .all(|(ln1, ln2): (&str, &str)| { + extract_indentation(ln1, checker.stylist) + == extract_indentation(ln2, checker.stylist) + }); + // Check if the code without the final return statement is only comments + if all_indentation_same + && code_is_only_comments( + in_between_source.trim_end().trim_end_matches("return"), + ) + { + // Avoid a trailing empty line and avoid causing invalid syntax with semicolon usage + let fixed_source = match in_between_source + .strip_prefix(checker.stylist.line_ending().as_str()) + { + Some(newline_stripped) => { + match newline_stripped + .strip_prefix(checker.stylist.indentation().as_str()) + { + Some(unindented) => unindented.to_owned(), + None => in_between_source, + } + } + None => match in_between_source.strip_prefix(';') { + Some(no_semicolon) => no_semicolon.to_owned(), + None => in_between_source, + }, + }; + diagnostic.amend(Fix::replacement( + fixed_source + &unparse_expr(assign_expr, checker.stylist), + last_assign_loc.location, + expr.end_location.unwrap(), + )); + } + } + } + } + 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 8b9969a243162..16e664bafab6d 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 @@ -10,7 +10,15 @@ expression: diagnostics end_location: row: 6 column: 12 - fix: ~ + fix: + content: + - return 1 + location: + row: 5 + column: 4 + end_location: + row: 6 + column: 12 parent: ~ - kind: UnnecessaryAssign: ~ @@ -20,7 +28,15 @@ expression: diagnostics end_location: row: 13 column: 12 - fix: ~ + fix: + content: + - return 2 + location: + row: 12 + column: 4 + end_location: + row: 13 + column: 12 parent: ~ - kind: UnnecessaryAssign: ~ @@ -35,21 +51,53 @@ expression: diagnostics - kind: UnnecessaryAssign: ~ location: - row: 31 + row: 33 column: 11 end_location: - row: 31 + row: 33 + column: 14 + fix: + content: + - "\"\"\"" + - " This is a" + - " block comment" + - " we want to keep" + - " \"\"\"" + - " # This is a comment that we want to keep." + - " return [5, 4, 3, 2, 1]" + location: + row: 26 + column: 4 + end_location: + row: 33 + column: 14 + parent: ~ +- kind: + UnnecessaryAssign: ~ + location: + row: 43 + column: 11 + end_location: + row: 43 column: 17 fix: ~ parent: ~ - kind: UnnecessaryAssign: ~ location: - row: 39 + row: 51 column: 11 end_location: - row: 39 + row: 51 column: 20 - fix: ~ + fix: + content: + - "return formatted.replace(\"()\", \"\").replace(\" \", \" \").strip()" + location: + row: 50 + column: 4 + end_location: + row: 51 + column: 20 parent: ~ diff --git a/crates/ruff/src/rules/flake8_return/visitor.rs b/crates/ruff/src/rules/flake8_return/visitor.rs index 74768eb91ec27..0f71886bd3157 100644 --- a/crates/ruff/src/rules/flake8_return/visitor.rs +++ b/crates/ruff/src/rules/flake8_return/visitor.rs @@ -1,6 +1,7 @@ use rustc_hash::{FxHashMap, FxHashSet}; use rustpython_parser::ast::{Expr, ExprKind, Location, Stmt, StmtKind}; +use crate::ast::types::Range; use crate::ast::visitor; use crate::ast::visitor::Visitor; @@ -12,9 +13,10 @@ pub struct Stack<'a> { pub elifs: Vec<&'a Stmt>, pub refs: FxHashMap<&'a str, Vec>, pub non_locals: FxHashSet<&'a str>, - pub assigns: FxHashMap<&'a str, Vec>, + pub assigns: FxHashMap<&'a str, Vec>, pub loops: Vec<(Location, Location)>, pub tries: Vec<(Location, Location)>, + pub assign_values: FxHashMap<&'a str, &'a Expr>, } #[derive(Default)] @@ -32,14 +34,6 @@ impl<'a> ReturnVisitor<'a> { } return; } - ExprKind::Name { id, .. } => { - self.stack - .assigns - .entry(id) - .or_insert_with(Vec::new) - .push(expr.location); - return; - } ExprKind::Attribute { .. } => { // Attribute assignments are often side-effects (e.g., `self.property = value`), // so we conservatively treat them as references to every known @@ -124,6 +118,21 @@ impl<'a> Visitor<'a> for ReturnVisitor<'a> { return; } + // TODO: Make this only run if the autofix is enabled (or we could merge into the below statement) + if let ExprKind::Name { id, .. } = &target.node { + self.stack.assign_values.insert(id, value); + } + + // This is hacky, but I don't kmow how to pass the "context" of the + // statement in another way. + if let ExprKind::Name { id, .. } = &target.node { + self.stack + .assigns + .entry(id) + .or_insert_with(Vec::new) + .push(Range::from_located(stmt)); + return; + } self.visit_assign_target(target); } }