Skip to content

Commit

Permalink
check for external references to the for loop variable
Browse files Browse the repository at this point in the history
  • Loading branch information
w0nder1ng committed Nov 18, 2024
1 parent f30118e commit 238d17a
Show file tree
Hide file tree
Showing 4 changed files with 78 additions and 1 deletion.
12 changes: 12 additions & 0 deletions crates/ruff_linter/resources/test/fixtures/perflint/PERF401.py
Original file line number Diff line number Diff line change
Expand Up @@ -156,3 +156,15 @@ def f():
result.append(1)
for i in range(10):
result.append(i*2) # PERF401

def f():
result = []
result += [1]
for i in range(10):
result.append(i*2) # PERF401

def f():
result = []
for val in range(5):
result.append(val * 2) # Ok
print(val)
Original file line number Diff line number Diff line change
Expand Up @@ -218,6 +218,34 @@ pub(crate) fn manual_list_comprehension(checker: &mut Checker, for_stmt: &ast::S
return;
}

// Avoid if the for-loop target is used outside of the for loop, e.g.,
//
// ```python
// for x in y:
// filtered.append(x)
// print(x)
// ```
//
// If this were a comprehension, x would no longer have the correct scope:
//
// ```python
// filtered = [x for x in y]
// print(x)
// ```
let for_loop_target = checker
.semantic()
.lookup_symbol(id.as_str())
.map(|id| checker.semantic().binding(id))
.expect("for loop target must exist");
// TODO: this currently does not properly find usages outside the for loop; figure out why
if for_loop_target
.references()
.map(|id| checker.semantic().reference(id))
.any(|reference| !for_stmt.range.contains_range(reference.range()))
{
return;
}

let binding_stmt = binding
.statement(checker.semantic())
.and_then(|stmt| stmt.as_assign_stmt());
Expand Down Expand Up @@ -261,7 +289,7 @@ pub(crate) fn manual_list_comprehension(checker: &mut Checker, for_stmt: &ast::S
// if there's more than one, then it's been accessed in the middle somewhere, so it's not safe to change into a list comprehension
count < 2
});

// If the binding has multiple statements on its line, it gets more complicated to fix, so we use an extend
// TODO: make this work
let binding_only_stmt_on_line = binding_stmt.is_some_and(|_binding_stmt| true);
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -121,5 +121,17 @@ PERF401.py:158:9: PERF401 Use `list.extend` to create a transformed list
157 | for i in range(10):
158 | result.append(i*2) # PERF401
| ^^^^^^^^^^^^^^^^^^ PERF401
159 |
160 | def f():
|
= help: Replace for loop with list.extend

PERF401.py:169:9: PERF401 Use a list comprehension to create a transformed list
|
167 | result = []
168 | for val in range(5):
169 | result.append(val * 2) # Ok
| ^^^^^^^^^^^^^^^^^^^^^^ PERF401
170 | print(val)
|
= help: Replace for loop with list comprehension
Original file line number Diff line number Diff line change
Expand Up @@ -279,6 +279,8 @@ PERF401.py:158:9: PERF401 [*] Use `list.extend` to create a transformed list
157 | for i in range(10):
158 | result.append(i*2) # PERF401
| ^^^^^^^^^^^^^^^^^^ PERF401
159 |
160 | def f():
|
= help: Replace for loop with list.extend

Expand All @@ -289,3 +291,26 @@ PERF401.py:158:9: PERF401 [*] Use `list.extend` to create a transformed list
157 |- for i in range(10):
158 |- result.append(i*2) # PERF401
157 |+ result.extend(i*2 for i in range(10)) # PERF401
159 158 |
160 159 | def f():
161 160 | result = []

PERF401.py:169:9: PERF401 [*] Use a list comprehension to create a transformed list
|
167 | result = []
168 | for val in range(5):
169 | result.append(val * 2) # Ok
| ^^^^^^^^^^^^^^^^^^^^^^ PERF401
170 | print(val)
|
= help: Replace for loop with list comprehension

Unsafe fix
164 164 | result.append(i*2) # PERF401
165 165 |
166 166 | def f():
167 |- result = []
168 |- for val in range(5):
169 |- result.append(val * 2) # Ok
167 |+ result = [val * 2 for val in range(5)] # Ok
170 168 | print(val)

0 comments on commit 238d17a

Please sign in to comment.