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

Co-authored-by: Micha Reiser <micha@reiser.io>
This commit is contained in:
w0nder1ng 2024-11-11 06:17:02 -05:00 committed by GitHub
parent 813ec23ecd
commit 5a3886c8b5
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
5 changed files with 491 additions and 10 deletions

View file

@ -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)

View file

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

View file

@ -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 {
"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<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,
}

View file

@ -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

View file

@ -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]