Skip to content

Commit

Permalink
Avoid PERF rules for iteration-dependent assignments (#5508)
Browse files Browse the repository at this point in the history
## 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 #5494.

Closes #5500.
  • Loading branch information
charliermarsh committed Jul 4, 2023
1 parent 75da72b commit c395e44
Show file tree
Hide file tree
Showing 4 changed files with 54 additions and 20 deletions.
7 changes: 7 additions & 0 deletions crates/ruff/resources/test/fixtures/perflint/PERF401.py
Original file line number Diff line number Diff line change
Expand Up @@ -30,3 +30,10 @@ def f():
result = []
for i in items:
result.append(i) # OK


def f():
items = [1, 2, 3, 4]
result = {}
for i in items:
result[i].append(i) # OK
7 changes: 7 additions & 0 deletions crates/ruff/resources/test/fixtures/perflint/PERF402.py
Original file line number Diff line number Diff line change
Expand Up @@ -17,3 +17,10 @@ def f():
result = []
for i in items:
result.append(i * i) # OK


def f():
items = [1, 2, 3, 4]
result = {}
for i in items:
result[i].append(i * i) # OK
30 changes: 20 additions & 10 deletions crates/ruff/src/rules/perflint/rules/manual_list_comprehension.rs
Original file line number Diff line number Diff line change
Expand Up @@ -2,6 +2,7 @@ use rustpython_parser::ast::{self, Expr, Stmt};

use ruff_diagnostics::{Diagnostic, Violation};
use ruff_macros::{derive_message_formats, violation};
use ruff_python_ast::helpers::any_over_expr;

use crate::checkers::ast::Checker;

Expand Down Expand Up @@ -52,6 +53,10 @@ impl Violation for ManualListComprehension {

/// PERF401
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 {
// ```python
// for x in y:
Expand Down Expand Up @@ -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)`).
if !conditional {
if arg.as_name_expr().map_or(false, |arg| {
target
.as_name_expr()
.map_or(false, |target| arg.id == target.id)
}) {
if arg.as_name_expr().map_or(false, |arg| arg.id == *id) {
return;
}
}

let Expr::Attribute(ast::ExprAttribute { attr, .. }) = func.as_ref() else {
let Expr::Attribute(ast::ExprAttribute { attr, value, .. }) = func.as_ref() else {
return;
};

if attr.as_str() == "append" {
checker
.diagnostics
.push(Diagnostic::new(ManualListComprehension, *range));
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
.diagnostics
.push(Diagnostic::new(ManualListComprehension, *range));
}
30 changes: 20 additions & 10 deletions crates/ruff/src/rules/perflint/rules/manual_list_copy.rs
Original file line number Diff line number Diff line change
Expand Up @@ -2,6 +2,7 @@ use rustpython_parser::ast::{self, Expr, Stmt};

use ruff_diagnostics::{Diagnostic, Violation};
use ruff_macros::{derive_message_formats, violation};
use ruff_python_ast::helpers::any_over_expr;

use crate::checkers::ast::Checker;

Expand Down Expand Up @@ -45,6 +46,10 @@ impl Violation for ManualListCopy {

/// PERF402
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 {
return;
};
Expand Down Expand Up @@ -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)`).
if !arg.as_name_expr().map_or(false, |arg| {
target
.as_name_expr()
.map_or(false, |target| arg.id == target.id)
}) {
if !arg.as_name_expr().map_or(false, |arg| arg.id == *id) {
return;
}

let Expr::Attribute(ast::ExprAttribute { attr, .. }) = func.as_ref() else {
let Expr::Attribute(ast::ExprAttribute { attr, value, .. }) = func.as_ref() else {
return;
};

if matches!(attr.as_str(), "append" | "insert") {
checker
.diagnostics
.push(Diagnostic::new(ManualListCopy, *range));
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;
}

checker
.diagnostics
.push(Diagnostic::new(ManualListCopy, *range));
}

0 comments on commit c395e44

Please sign in to comment.