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

crates/ruff_linter/resources/test/fixtures/perflint/PERF401.py

@@ -87,3 +87,16 @@ 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):
+            if i % 2:
+                result.append(i) # PERF401

crates/ruff_linter/src/rules/perflint/rules/manual_list_comprehension.rs

@@ -1,10 +1,13 @@
 use ruff_python_ast::{self as ast, Arguments, Expr, Stmt};
 
-use ruff_diagnostics::{Diagnostic, Violation};
+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_text_size::TextRange;
+
+use anyhow::{anyhow, Result};
 
 use crate::checkers::ast::Checker;
 
@@ -46,17 +49,30 @@ 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 {
-        let ManualListComprehension { is_async } = self;
+        let ManualListComprehension { is_async, .. } = self;
         match is_async {
             false => format!("Use a list comprehension to create a transformed list"),
             true => format!("Use an async list comprehension to create a transformed list"),
         }
     }
+
+    fn fix_title(&self) -> Option<String> {
+        self.comprehension_type
+            .map(|comprehension_type| match comprehension_type {
+                ComprehensionType::ListComprehension => {
+                    format!("Replace for loop with list comprehension")
+                }
+                ComprehensionType::Extend => format!("Replace for loop with list.extend"),
+            })
+    }
 }
 
 /// PERF401
@@ -126,7 +142,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 +203,159 @@ 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,
+        });
+
+    let comprehension_type = if binding_is_empty_list {
+        ComprehensionType::ListComprehension
+    } else {
+        ComprehensionType::Extend
+    };
+
+    // 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 = {
+        let for_loop_parent = checker.semantic().current_statement_parent_id();
+
+        binding.source.is_some_and(|binding_source| {
+            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 = {
+        let only_target = binding_stmt.is_some_and(|binding_stmt| binding_stmt.targets.len() == 1);
+        let is_name =
+            binding_stmt.is_some_and(|binding_stmt| binding_stmt.targets[0].is_name_expr());
+        only_target && is_name
+    };
+
+    // A list extend works in every context, while a list comprehension only works when all the criteria are true
+    let comprehension_type =
+        Some(comprehension_type).filter(|comprehension_type| match comprehension_type {
+            ComprehensionType::ListComprehension => {
+                binding_stmt.is_some() && assignment_in_same_statement && binding_has_one_target
+            }
+            ComprehensionType::Extend => true,
+        });
+
+    let mut diagnostic = Diagnostic::new(
         ManualListComprehension {
             is_async: for_stmt.is_async,
+            comprehension_type,
         },
         *range,
-    ));
+    );
+
+    diagnostic.try_set_optional_fix(|| match comprehension_type {
+        Some(comprehension_type) => convert_to_list_extend(
+            comprehension_type,
+            binding,
+            for_stmt,
+            if_test.map(std::convert::AsRef::as_ref),
+            arg,
+            checker,
+        )
+        .map(Some),
+        None => Ok(None),
+    });
+
+    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 comprehension = ast::Comprehension {
+        target: (*for_stmt.target).clone(),
+        iter: (*for_stmt.iter).clone(),
+        is_async: for_stmt.is_async,
+        ifs: if_test.into_iter().cloned().collect(),
+        range: TextRange::default(),
+    };
+    match fix_type {
+        ComprehensionType::Extend => {
+            let generator = ast::ExprGenerator {
+                elt: Box::new(to_append.clone()),
+                generators: vec![comprehension],
+                parenthesized: false,
+                range: TextRange::default(),
+            };
+
+            let variable_name = checker.locator().slice(binding.range);
+
+            let extend = ast::ExprAttribute {
+                value: Box::new(Expr::Name(ast::ExprName {
+                    id: ast::name::Name::new(variable_name),
+                    ctx: ast::ExprContext::Load,
+                    range: TextRange::default(),
+                })),
+                attr: ast::Identifier::new("extend", TextRange::default()),
+                ctx: ast::ExprContext::Load,
+                range: TextRange::default(),
+            };
+
+            let list_extend = ast::ExprCall {
+                func: Box::new(ast::Expr::Attribute(extend)),
+                arguments: Arguments {
+                    args: Box::new([ast::Expr::Generator(generator)]),
+                    range: TextRange::default(),
+                    keywords: Box::new([]),
+                },
+                range: TextRange::default(),
+            };
+
+            let comprehension_body = checker.generator().expr(&Expr::Call(list_extend));
+
+            Ok(Fix::unsafe_edit(Edit::range_replacement(
+                comprehension_body,
+                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 list_comp = ast::ExprListComp {
+                elt: Box::new(to_append.clone()),
+                generators: vec![comprehension],
+                range: TextRange::default(),
+            };
+
+            let comprehension_body = checker.generator().expr(&Expr::ListComp(list_comp));
+            Ok(Fix::unsafe_edits(
+                Edit::range_replacement(comprehension_body, empty_list_to_replace.range),
+                [Edit::range_deletion(for_stmt.range)],
+            ))
+        }
+    }
+}
+
+#[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

@@ -1,34 +1,121 @@
 ---
 source: crates/ruff_linter/src/rules/perflint/mod.rs
 ---
-PERF401.py:6:13: PERF401 Use a list comprehension to create a transformed list
+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
 
-PERF401.py:13:9: PERF401 Use a list comprehension to create a transformed list
+ℹ 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 |+    result = [i for i in items if i % 2]
+  4 |+      # PERF401
+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 |+    result = [i * i for i in items]
+   12 |+      # PERF401
+14 13 | 
+15 14 | 
+16 15 | def f():
 
-PERF401.py:82:13: PERF401 Use an async list comprehension to create a transformed list
+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 |+    result = [i async for i in items if i % 2]
+   80 |+      # PERF401
+83 81 | 
+84 82 | 
+85 83 | def f():
 
-PERF401.py:89:9: PERF401 Use an async list comprehension to create a transformed list
+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 |+    result = [i async for i in items]
+   88 |+      # PERF401
+90 89 | 
+91 90 | 
+92 91 | def f():
+
+PERF401.py:95:9: PERF401 [*] Use a list comprehension to create a transformed list
+   |
+93 |     result, _ = [1,2,3,4], ...
+94 |     for i in range(10):
+95 |         result.append(i*2)  # PERF401
+   |         ^^^^^^^^^^^^^^^^^^ PERF401
+96 | 
+97 | def f():
+   |
+   = 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 | def f():
+98 97 |     result = []
+
+PERF401.py:102:17: PERF401 Use a list comprehension to create a transformed list
+    |
+100 |         for i in range(10):
+101 |             if i % 2:
+102 |                 result.append(i) # PERF401
+    |                 ^^^^^^^^^^^^^^^^ PERF401
+    |