[perflint] implement quick-fix for manual-list-comprehension (PERF401) (#13919)

Co-authored-by: Micha Reiser <micha@reiser.io>

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)

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(())
+    }
 }

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<ComprehensionType>,
 }
 
 impl Violation for ManualListComprehension {
+    const FIX_AVAILABILITY: FixAvailability = FixAvailability::Sometimes;
+
     #[derive_message_formats]
     fn message(&self) -> String {
-        if self.is_async {
+        let ManualListComprehension {
-            "Use an async list comprehension to create a transformed list".to_string()
+            is_async,
-        } else {
+            comprehension_type,
-            "Use a list comprehension to create a transformed list".to_string()
+        } = 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<String> {
+        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<Fix> {
+    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,
 }

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

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]