Avoid PERF rules for iteration-dependent assignments (#5508)

## Summary

We need to avoid raising "rewrite as a comprehension" violations in
cases like:

```python
d = defaultdict(list)

for i in [1, 2, 3]:
    d[i].append(i**2)
```

Closes https://github.com/astral-sh/ruff/issues/5494.

Closes https://github.com/astral-sh/ruff/issues/5500.
This commit is contained in:
Charlie Marsh 2023-07-04 14:21:05 -04:00 committed by GitHub
parent 75da72bd7f
commit c395e44bd7
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
4 changed files with 57 additions and 23 deletions

View File

@ -30,3 +30,10 @@ def f():
result = [] result = []
for i in items: for i in items:
result.append(i) # OK result.append(i) # OK
def f():
items = [1, 2, 3, 4]
result = {}
for i in items:
result[i].append(i) # OK

View File

@ -17,3 +17,10 @@ def f():
result = [] result = []
for i in items: for i in items:
result.append(i * i) # OK result.append(i * i) # OK
def f():
items = [1, 2, 3, 4]
result = {}
for i in items:
result[i].append(i * i) # OK

View File

@ -2,6 +2,7 @@ use rustpython_parser::ast::{self, Expr, Stmt};
use ruff_diagnostics::{Diagnostic, Violation}; use ruff_diagnostics::{Diagnostic, Violation};
use ruff_macros::{derive_message_formats, violation}; use ruff_macros::{derive_message_formats, violation};
use ruff_python_ast::helpers::any_over_expr;
use crate::checkers::ast::Checker; use crate::checkers::ast::Checker;
@ -52,6 +53,10 @@ impl Violation for ManualListComprehension {
/// PERF401 /// PERF401
pub(crate) fn manual_list_comprehension(checker: &mut Checker, target: &Expr, body: &[Stmt]) { pub(crate) fn manual_list_comprehension(checker: &mut Checker, target: &Expr, body: &[Stmt]) {
let Expr::Name(ast::ExprName { id, .. }) = target else {
return;
};
let (stmt, conditional) = match body { let (stmt, conditional) = match body {
// ```python // ```python
// for x in y: // for x in y:
@ -99,22 +104,27 @@ pub(crate) fn manual_list_comprehension(checker: &mut Checker, target: &Expr, bo
// Ignore direct list copies (e.g., `for x in y: filtered.append(x)`). // Ignore direct list copies (e.g., `for x in y: filtered.append(x)`).
if !conditional { if !conditional {
if arg.as_name_expr().map_or(false, |arg| { if arg.as_name_expr().map_or(false, |arg| arg.id == *id) {
target
.as_name_expr()
.map_or(false, |target| arg.id == target.id)
}) {
return; return;
} }
} }
let Expr::Attribute(ast::ExprAttribute { attr, .. }) = func.as_ref() else { let Expr::Attribute(ast::ExprAttribute { attr, value, .. }) = func.as_ref() else {
return; return;
}; };
if attr.as_str() == "append" { if attr.as_str() != "append" {
return;
}
// Avoid, e.g., `for x in y: filtered[x].append(x * x)`.
if any_over_expr(value, &|expr| {
expr.as_name_expr().map_or(false, |expr| expr.id == *id)
}) {
return;
}
checker checker
.diagnostics .diagnostics
.push(Diagnostic::new(ManualListComprehension, *range)); .push(Diagnostic::new(ManualListComprehension, *range));
}
} }

View File

@ -2,6 +2,7 @@ use rustpython_parser::ast::{self, Expr, Stmt};
use ruff_diagnostics::{Diagnostic, Violation}; use ruff_diagnostics::{Diagnostic, Violation};
use ruff_macros::{derive_message_formats, violation}; use ruff_macros::{derive_message_formats, violation};
use ruff_python_ast::helpers::any_over_expr;
use crate::checkers::ast::Checker; use crate::checkers::ast::Checker;
@ -45,6 +46,10 @@ impl Violation for ManualListCopy {
/// PERF402 /// PERF402
pub(crate) fn manual_list_copy(checker: &mut Checker, target: &Expr, body: &[Stmt]) { pub(crate) fn manual_list_copy(checker: &mut Checker, target: &Expr, body: &[Stmt]) {
let Expr::Name(ast::ExprName { id, .. }) = target else {
return;
};
let [stmt] = body else { let [stmt] = body else {
return; return;
}; };
@ -72,21 +77,26 @@ pub(crate) fn manual_list_copy(checker: &mut Checker, target: &Expr, body: &[Stm
}; };
// Only flag direct list copies (e.g., `for x in y: filtered.append(x)`). // Only flag direct list copies (e.g., `for x in y: filtered.append(x)`).
if !arg.as_name_expr().map_or(false, |arg| { if !arg.as_name_expr().map_or(false, |arg| arg.id == *id) {
target return;
.as_name_expr() }
.map_or(false, |target| arg.id == target.id)
let Expr::Attribute(ast::ExprAttribute { attr, value, .. }) = func.as_ref() else {
return;
};
if !matches!(attr.as_str(), "append" | "insert") {
return;
}
// Avoid, e.g., `for x in y: filtered[x].append(x * x)`.
if any_over_expr(value, &|expr| {
expr.as_name_expr().map_or(false, |expr| expr.id == *id)
}) { }) {
return; return;
} }
let Expr::Attribute(ast::ExprAttribute { attr, .. }) = func.as_ref() else {
return;
};
if matches!(attr.as_str(), "append" | "insert") {
checker checker
.diagnostics .diagnostics
.push(Diagnostic::new(ManualListCopy, *range)); .push(Diagnostic::new(ManualListCopy, *range));
}
} }