diff --git a/crates/ruff_linter/resources/test/fixtures/perflint/PERF401.py b/crates/ruff_linter/resources/test/fixtures/perflint/PERF401.py index f8b8014745017..bead4d5ce0afa 100644 --- a/crates/ruff_linter/resources/test/fixtures/perflint/PERF401.py +++ b/crates/ruff_linter/resources/test/fixtures/perflint/PERF401.py @@ -87,3 +87,41 @@ def f(): result = [] async for i in items: result.append(i) # PERF401 + + +def f(): + result, _ = [1,2,3,4], ... + for i in range(10): + result.append(i*2) # PERF401 + + +def f(): + result = [] + if True: + for i in range(10): # single-line comment 1 should be protected + # single-line comment 2 should be protected + if i % 2: # single-line comment 3 should be protected + result.append(i) # PERF401 + + +def f(): + result = [] # comment after assignment should be protected + for i in range(10): # single-line comment 1 should be protected + # single-line comment 2 should be protected + if i % 2: # single-line comment 3 should be protected + result.append(i) # PERF401 + + +def f(): + result = [] + for i in range(10): + """block comment stops the fix""" + result.append(i*2) # Ok + +def f(param): + # PERF401 + # make sure the fix does not panic if there is no comments + if param: + new_layers = [] + for value in param: + new_layers.append(value * 3) diff --git a/crates/ruff_linter/src/rules/perflint/mod.rs b/crates/ruff_linter/src/rules/perflint/mod.rs index 2d073e639e242..55477f8d1f924 100644 --- a/crates/ruff_linter/src/rules/perflint/mod.rs +++ b/crates/ruff_linter/src/rules/perflint/mod.rs @@ -10,7 +10,7 @@ mod tests { use crate::assert_messages; use crate::registry::Rule; - use crate::settings::types::PythonVersion; + use crate::settings::types::{PreviewMode, PythonVersion}; use crate::settings::LinterSettings; use crate::test::test_path; @@ -29,4 +29,24 @@ mod tests { assert_messages!(snapshot, diagnostics); Ok(()) } + + // TODO: remove this test case when the fix for `perf401` is stabilized + #[test_case(Rule::ManualListComprehension, Path::new("PERF401.py"))] + fn preview_rules(rule_code: Rule, path: &Path) -> Result<()> { + let snapshot = format!( + "preview__{}_{}", + rule_code.noqa_code(), + path.to_string_lossy() + ); + let diagnostics = test_path( + Path::new("perflint").join(path).as_path(), + &LinterSettings { + preview: PreviewMode::Enabled, + target_version: PythonVersion::Py310, + ..LinterSettings::for_rule(rule_code) + }, + )?; + assert_messages!(snapshot, diagnostics); + Ok(()) + } } diff --git a/crates/ruff_linter/src/rules/perflint/rules/manual_list_comprehension.rs b/crates/ruff_linter/src/rules/perflint/rules/manual_list_comprehension.rs index b65c403f29f7d..e9c3a35da0da0 100644 --- a/crates/ruff_linter/src/rules/perflint/rules/manual_list_comprehension.rs +++ b/crates/ruff_linter/src/rules/perflint/rules/manual_list_comprehension.rs @@ -1,10 +1,14 @@ use ruff_python_ast::{self as ast, Arguments, Expr, Stmt}; -use ruff_diagnostics::{Diagnostic, Violation}; +use anyhow::{anyhow, Result}; +use ruff_diagnostics::{Diagnostic, Edit, Fix, FixAvailability, Violation}; use ruff_macros::{derive_message_formats, violation}; use ruff_python_ast::comparable::ComparableExpr; use ruff_python_ast::helpers::any_over_expr; -use ruff_python_semantic::analyze::typing::is_list; +use ruff_python_semantic::{analyze::typing::is_list, Binding}; +use ruff_python_trivia::PythonWhitespace; +use ruff_source_file::LineRanges; +use ruff_text_size::{Ranged, TextRange}; use crate::checkers::ast::Checker; @@ -46,15 +50,43 @@ use crate::checkers::ast::Checker; #[violation] pub struct ManualListComprehension { is_async: bool, + comprehension_type: Option, } impl Violation for ManualListComprehension { + const FIX_AVAILABILITY: FixAvailability = FixAvailability::Sometimes; + #[derive_message_formats] fn message(&self) -> String { - if self.is_async { - "Use an async list comprehension to create a transformed list".to_string() - } else { - "Use a list comprehension to create a transformed list".to_string() + let ManualListComprehension { + is_async, + comprehension_type, + } = self; + let message_str = match comprehension_type { + Some(ComprehensionType::Extend) => { + if *is_async { + "`list.extend` with an async comprehension" + } else { + "`list.extend`" + } + } + Some(ComprehensionType::ListComprehension) | None => { + if *is_async { + "an async list comprehension" + } else { + "a list comprehension" + } + } + }; + format!("Use {message_str} to create a transformed list") + } + + fn fix_title(&self) -> Option { + match self.comprehension_type? { + ComprehensionType::ListComprehension => { + Some("Replace for loop with list comprehension".to_string()) + } + ComprehensionType::Extend => Some("Replace for loop with list.extend".to_string()), } } } @@ -126,7 +158,6 @@ pub(crate) fn manual_list_comprehension(checker: &mut Checker, for_stmt: &ast::S if attr.as_str() != "append" { return; } - // Ignore direct list copies (e.g., `for x in y: filtered.append(x)`), unless it's async, which // `manual-list-copy` doesn't cover. if !for_stmt.is_async { @@ -188,10 +219,180 @@ pub(crate) fn manual_list_comprehension(checker: &mut Checker, for_stmt: &ast::S return; } - checker.diagnostics.push(Diagnostic::new( + let binding_stmt = binding + .statement(checker.semantic()) + .and_then(|stmt| stmt.as_assign_stmt()); + + // If the variable is an empty list literal, then we might be able to replace it with a full list comprehension + // otherwise, it has to be replaced with a `list.extend` + let binding_is_empty_list = + binding_stmt.is_some_and(|binding_stmt| match binding_stmt.value.as_list_expr() { + Some(list_expr) => list_expr.elts.is_empty(), + None => false, + }); + + // If the for loop does not have the same parent element as the binding, then it cannot always be + // deleted and replaced with a list comprehension. This does not apply when using an extend. + let assignment_in_same_statement = { + binding.source.is_some_and(|binding_source| { + let for_loop_parent = checker.semantic().current_statement_parent_id(); + let binding_parent = checker.semantic().parent_statement_id(binding_source); + for_loop_parent == binding_parent + }) + }; + + // If the binding is not a single name expression, it could be replaced with a list comprehension, + // but not necessarily, so this needs to be manually fixed. This does not apply when using an extend. + let binding_has_one_target = { + match binding_stmt.map(|binding_stmt| binding_stmt.targets.as_slice()) { + Some([only_target]) => only_target.is_name_expr(), + _ => false, + } + }; + + // A list extend works in every context, while a list comprehension only works when all the criteria are true + let comprehension_type = + if binding_is_empty_list && assignment_in_same_statement && binding_has_one_target { + ComprehensionType::ListComprehension + } else { + ComprehensionType::Extend + }; + + let mut diagnostic = Diagnostic::new( ManualListComprehension { is_async: for_stmt.is_async, + comprehension_type: Some(comprehension_type), }, *range, - )); + ); + + // TODO: once this fix is stabilized, change the rule to always fixable + if checker.settings.preview.is_enabled() { + diagnostic.try_set_fix(|| { + convert_to_list_extend( + comprehension_type, + binding, + for_stmt, + if_test.map(std::convert::AsRef::as_ref), + arg, + checker, + ) + }); + } + + checker.diagnostics.push(diagnostic); +} + +fn convert_to_list_extend( + fix_type: ComprehensionType, + binding: &Binding, + for_stmt: &ast::StmtFor, + if_test: Option<&ast::Expr>, + to_append: &Expr, + checker: &Checker, +) -> Result { + let locator = checker.locator(); + let if_str = match if_test { + Some(test) => format!(" if {}", locator.slice(test.range())), + None => String::new(), + }; + + let for_iter_str = locator.slice(for_stmt.iter.range()); + let for_type = if for_stmt.is_async { + "async for" + } else { + "for" + }; + let target_str = locator.slice(for_stmt.target.range()); + let elt_str = locator.slice(to_append); + let generator_str = format!("{elt_str} {for_type} {target_str} in {for_iter_str}{if_str}"); + + let comment_strings_in_range = |range| { + checker + .comment_ranges() + .comments_in_range(range) + .iter() + .map(|range| locator.slice(range).trim_whitespace_start()) + .collect() + }; + let for_stmt_end = for_stmt.range.end(); + let for_loop_inline_comments: Vec<&str> = comment_strings_in_range(for_stmt.range); + let for_loop_trailing_comment = + comment_strings_in_range(TextRange::new(for_stmt_end, locator.line_end(for_stmt_end))); + let newline = checker.stylist().line_ending().as_str(); + + match fix_type { + ComprehensionType::Extend => { + let variable_name = checker.locator().slice(binding.range); + + let comprehension_body = format!("{variable_name}.extend({generator_str})"); + + let indent_range = TextRange::new( + locator.line_start(for_stmt.range.start()), + for_stmt.range.start(), + ); + let indentation = if for_loop_inline_comments.is_empty() { + String::new() + } else { + format!("{newline}{}", locator.slice(indent_range)) + }; + let text_to_replace = format!( + "{}{indentation}{comprehension_body}", + for_loop_inline_comments.join(&indentation) + ); + Ok(Fix::unsafe_edit(Edit::range_replacement( + text_to_replace, + for_stmt.range, + ))) + } + ComprehensionType::ListComprehension => { + let binding_stmt = binding + .statement(checker.semantic()) + .and_then(|stmt| stmt.as_assign_stmt()) + .ok_or(anyhow!( + "Binding must have a statement to convert into a list comprehension" + ))?; + let empty_list_to_replace = binding_stmt.value.as_list_expr().ok_or(anyhow!( + "Assignment value must be an empty list literal in order to replace with a list comprehension" + ))?; + + let comprehension_body = format!("[{generator_str}]"); + + let indent_range = TextRange::new( + locator.line_start(binding_stmt.range.start()), + binding_stmt.range.start(), + ); + + let mut for_loop_comments = for_loop_inline_comments; + for_loop_comments.extend(for_loop_trailing_comment); + + let indentation = if for_loop_comments.is_empty() { + String::new() + } else { + format!("{newline}{}", locator.slice(indent_range)) + }; + let leading_comments = format!("{}{indentation}", for_loop_comments.join(&indentation)); + + let mut additional_fixes = vec![Edit::range_deletion( + locator.full_lines_range(for_stmt.range), + )]; + // if comments are empty, trying to insert them panics + if !leading_comments.is_empty() { + additional_fixes.push(Edit::insertion( + leading_comments, + binding_stmt.range.start(), + )); + } + Ok(Fix::unsafe_edits( + Edit::range_replacement(comprehension_body, empty_list_to_replace.range), + additional_fixes, + )) + } + } +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +enum ComprehensionType { + Extend, + ListComprehension, } diff --git a/crates/ruff_linter/src/rules/perflint/snapshots/ruff_linter__rules__perflint__tests__PERF401_PERF401.py.snap b/crates/ruff_linter/src/rules/perflint/snapshots/ruff_linter__rules__perflint__tests__PERF401_PERF401.py.snap index ae0dcb8e5a10f..31cf524b6d3a7 100644 --- a/crates/ruff_linter/src/rules/perflint/snapshots/ruff_linter__rules__perflint__tests__PERF401_PERF401.py.snap +++ b/crates/ruff_linter/src/rules/perflint/snapshots/ruff_linter__rules__perflint__tests__PERF401_PERF401.py.snap @@ -8,6 +8,7 @@ PERF401.py:6:13: PERF401 Use a list comprehension to create a transformed list 6 | result.append(i) # PERF401 | ^^^^^^^^^^^^^^^^ PERF401 | + = help: Replace for loop with list comprehension PERF401.py:13:9: PERF401 Use a list comprehension to create a transformed list | @@ -16,6 +17,7 @@ PERF401.py:13:9: PERF401 Use a list comprehension to create a transformed list 13 | result.append(i * i) # PERF401 | ^^^^^^^^^^^^^^^^^^^^ PERF401 | + = help: Replace for loop with list comprehension PERF401.py:82:13: PERF401 Use an async list comprehension to create a transformed list | @@ -24,6 +26,7 @@ PERF401.py:82:13: PERF401 Use an async list comprehension to create a transforme 82 | result.append(i) # PERF401 | ^^^^^^^^^^^^^^^^ PERF401 | + = help: Replace for loop with list comprehension PERF401.py:89:9: PERF401 Use an async list comprehension to create a transformed list | @@ -32,3 +35,40 @@ PERF401.py:89:9: PERF401 Use an async list comprehension to create a transformed 89 | result.append(i) # PERF401 | ^^^^^^^^^^^^^^^^ PERF401 | + = help: Replace for loop with list comprehension + +PERF401.py:95:9: PERF401 Use `list.extend` to create a transformed list + | +93 | result, _ = [1,2,3,4], ... +94 | for i in range(10): +95 | result.append(i*2) # PERF401 + | ^^^^^^^^^^^^^^^^^^ PERF401 + | + = help: Replace for loop with list.extend + +PERF401.py:104:17: PERF401 Use `list.extend` to create a transformed list + | +102 | # single-line comment 2 should be protected +103 | if i % 2: # single-line comment 3 should be protected +104 | result.append(i) # PERF401 + | ^^^^^^^^^^^^^^^^ PERF401 + | + = help: Replace for loop with list.extend + +PERF401.py:112:13: PERF401 Use a list comprehension to create a transformed list + | +110 | # single-line comment 2 should be protected +111 | if i % 2: # single-line comment 3 should be protected +112 | result.append(i) # PERF401 + | ^^^^^^^^^^^^^^^^ PERF401 + | + = help: Replace for loop with list comprehension + +PERF401.py:127:13: PERF401 Use a list comprehension to create a transformed list + | +125 | new_layers = [] +126 | for value in param: +127 | new_layers.append(value * 3) + | ^^^^^^^^^^^^^^^^^^^^^^^^^^^^ PERF401 + | + = help: Replace for loop with list comprehension diff --git a/crates/ruff_linter/src/rules/perflint/snapshots/ruff_linter__rules__perflint__tests__preview__PERF401_PERF401.py.snap b/crates/ruff_linter/src/rules/perflint/snapshots/ruff_linter__rules__perflint__tests__preview__PERF401_PERF401.py.snap new file mode 100644 index 0000000000000..a981285c8b4a5 --- /dev/null +++ b/crates/ruff_linter/src/rules/perflint/snapshots/ruff_linter__rules__perflint__tests__preview__PERF401_PERF401.py.snap @@ -0,0 +1,182 @@ +--- +source: crates/ruff_linter/src/rules/perflint/mod.rs +snapshot_kind: text +--- +PERF401.py:6:13: PERF401 [*] Use a list comprehension to create a transformed list + | +4 | for i in items: +5 | if i % 2: +6 | result.append(i) # PERF401 + | ^^^^^^^^^^^^^^^^ PERF401 + | + = help: Replace for loop with list comprehension + +ℹ Unsafe fix +1 1 | def f(): +2 2 | items = [1, 2, 3, 4] +3 |- result = [] +4 |- for i in items: +5 |- if i % 2: +6 |- result.append(i) # PERF401 + 3 |+ # PERF401 + 4 |+ result = [i for i in items if i % 2] +7 5 | +8 6 | +9 7 | def f(): + +PERF401.py:13:9: PERF401 [*] Use a list comprehension to create a transformed list + | +11 | result = [] +12 | for i in items: +13 | result.append(i * i) # PERF401 + | ^^^^^^^^^^^^^^^^^^^^ PERF401 + | + = help: Replace for loop with list comprehension + +ℹ Unsafe fix +8 8 | +9 9 | def f(): +10 10 | items = [1, 2, 3, 4] +11 |- result = [] +12 |- for i in items: +13 |- result.append(i * i) # PERF401 + 11 |+ # PERF401 + 12 |+ result = [i * i for i in items] +14 13 | +15 14 | +16 15 | def f(): + +PERF401.py:82:13: PERF401 [*] Use an async list comprehension to create a transformed list + | +80 | async for i in items: +81 | if i % 2: +82 | result.append(i) # PERF401 + | ^^^^^^^^^^^^^^^^ PERF401 + | + = help: Replace for loop with list comprehension + +ℹ Unsafe fix +76 76 | +77 77 | def f(): +78 78 | items = [1, 2, 3, 4] +79 |- result = [] +80 |- async for i in items: +81 |- if i % 2: +82 |- result.append(i) # PERF401 + 79 |+ # PERF401 + 80 |+ result = [i async for i in items if i % 2] +83 81 | +84 82 | +85 83 | def f(): + +PERF401.py:89:9: PERF401 [*] Use an async list comprehension to create a transformed list + | +87 | result = [] +88 | async for i in items: +89 | result.append(i) # PERF401 + | ^^^^^^^^^^^^^^^^ PERF401 + | + = help: Replace for loop with list comprehension + +ℹ Unsafe fix +84 84 | +85 85 | def f(): +86 86 | items = [1, 2, 3, 4] +87 |- result = [] +88 |- async for i in items: +89 |- result.append(i) # PERF401 + 87 |+ # PERF401 + 88 |+ result = [i async for i in items] +90 89 | +91 90 | +92 91 | def f(): + +PERF401.py:95:9: PERF401 [*] Use `list.extend` to create a transformed list + | +93 | result, _ = [1,2,3,4], ... +94 | for i in range(10): +95 | result.append(i*2) # PERF401 + | ^^^^^^^^^^^^^^^^^^ PERF401 + | + = help: Replace for loop with list.extend + +ℹ Unsafe fix +91 91 | +92 92 | def f(): +93 93 | result, _ = [1,2,3,4], ... +94 |- for i in range(10): +95 |- result.append(i*2) # PERF401 + 94 |+ result.extend(i*2 for i in range(10)) # PERF401 +96 95 | +97 96 | +98 97 | def f(): + +PERF401.py:104:17: PERF401 [*] Use `list.extend` to create a transformed list + | +102 | # single-line comment 2 should be protected +103 | if i % 2: # single-line comment 3 should be protected +104 | result.append(i) # PERF401 + | ^^^^^^^^^^^^^^^^ PERF401 + | + = help: Replace for loop with list.extend + +ℹ Unsafe fix +98 98 | def f(): +99 99 | result = [] +100 100 | if True: +101 |- for i in range(10): # single-line comment 1 should be protected +102 |- # single-line comment 2 should be protected +103 |- if i % 2: # single-line comment 3 should be protected +104 |- result.append(i) # PERF401 + 101 |+ # single-line comment 1 should be protected + 102 |+ # single-line comment 2 should be protected + 103 |+ # single-line comment 3 should be protected + 104 |+ result.extend(i for i in range(10) if i % 2) # PERF401 +105 105 | +106 106 | +107 107 | def f(): + +PERF401.py:112:13: PERF401 [*] Use a list comprehension to create a transformed list + | +110 | # single-line comment 2 should be protected +111 | if i % 2: # single-line comment 3 should be protected +112 | result.append(i) # PERF401 + | ^^^^^^^^^^^^^^^^ PERF401 + | + = help: Replace for loop with list comprehension + +ℹ Unsafe fix +105 105 | +106 106 | +107 107 | def f(): +108 |- result = [] # comment after assignment should be protected +109 |- for i in range(10): # single-line comment 1 should be protected +110 |- # single-line comment 2 should be protected +111 |- if i % 2: # single-line comment 3 should be protected +112 |- result.append(i) # PERF401 + 108 |+ # single-line comment 1 should be protected + 109 |+ # single-line comment 2 should be protected + 110 |+ # single-line comment 3 should be protected + 111 |+ # PERF401 + 112 |+ result = [i for i in range(10) if i % 2] # comment after assignment should be protected +113 113 | +114 114 | +115 115 | def f(): + +PERF401.py:127:13: PERF401 [*] Use a list comprehension to create a transformed list + | +125 | new_layers = [] +126 | for value in param: +127 | new_layers.append(value * 3) + | ^^^^^^^^^^^^^^^^^^^^^^^^^^^^ PERF401 + | + = help: Replace for loop with list comprehension + +ℹ Unsafe fix +122 122 | # PERF401 +123 123 | # make sure the fix does not panic if there is no comments +124 124 | if param: +125 |- new_layers = [] +126 |- for value in param: +127 |- new_layers.append(value * 3) + 125 |+ new_layers = [value * 3 for value in param]