Skip to content

Commit

Permalink
[ruff] Expand rule for list(iterable).pop(0) idiom (RUF015) (as…
Browse files Browse the repository at this point in the history
…tral-sh#10148)

## Summary

Currently, rule `RUF015` is not able to detect the usage of
`list(iterable).pop(0)` falling under the category of an _unnecessary
iterable allocation for accessing the first element_. This PR wants to
change that. See the underlying issue for more details.

* Provide extension to detect `list(iterable).pop(0)`, but not
`list(iterable).pop(i)` where i > 1
* Update corresponding doc

## Test Plan

* `RUF015.py` and the corresponding snap file were extended such that
their correspond to the new behaviour

Closes astral-sh#9190

--- 

PS: I've only been working on this ticket as I haven't seen any activity
from issue assignee @rmad17, neither in this repo nor in a fork. I hope
I interpreted his inactivity correctly. Didn't mean to steal his chance.
Since I stumbled across the underlying problem myself, I wanted to offer
a solution as soon as possible.
  • Loading branch information
robincaloudis authored and nkxxll committed Mar 4, 2024
1 parent b9d8b2d commit 0851d42
Show file tree
Hide file tree
Showing 4 changed files with 146 additions and 42 deletions.
10 changes: 10 additions & 0 deletions crates/ruff_linter/resources/test/fixtures/ruff/RUF015.py
Expand Up @@ -57,6 +57,16 @@
list(zip(x, y))[0]
[*zip(x, y)][0]

# RUF015 (pop)
list(x).pop(0)
[i for i in x].pop(0)
list(i for i in x).pop(0)

# OK
list(x).pop(1)
list(x).remove(0)
list(x).remove(1)


def test():
zip = list # Overwrite the builtin zip
Expand Down
5 changes: 4 additions & 1 deletion crates/ruff_linter/src/checkers/ast/analyze/expression.rs
Expand Up @@ -116,7 +116,7 @@ pub(crate) fn expression(expr: &Expr, checker: &mut Checker) {
flake8_simplify::rules::use_capital_environment_variables(checker, expr);
}
if checker.enabled(Rule::UnnecessaryIterableAllocationForFirstElement) {
ruff::rules::unnecessary_iterable_allocation_for_first_element(checker, subscript);
ruff::rules::unnecessary_iterable_allocation_for_first_element(checker, expr);
}
if checker.enabled(Rule::InvalidIndexType) {
ruff::rules::invalid_index_type(checker, subscript);
Expand Down Expand Up @@ -965,6 +965,9 @@ pub(crate) fn expression(expr: &Expr, checker: &mut Checker) {
if checker.enabled(Rule::DefaultFactoryKwarg) {
ruff::rules::default_factory_kwarg(checker, call);
}
if checker.enabled(Rule::UnnecessaryIterableAllocationForFirstElement) {
ruff::rules::unnecessary_iterable_allocation_for_first_element(checker, expr);
}
}
Expr::Dict(dict) => {
if checker.any_enabled(&[
Expand Down
Expand Up @@ -11,14 +11,21 @@ use crate::checkers::ast::Checker;
use crate::fix::snippet::SourceCodeSnippet;

/// ## What it does
/// Checks for uses of `list(...)[0]` that can be replaced with
/// `next(iter(...))`.
/// Checks the following constructs, all of which can be replaced by
/// `next(iter(...))`:
///
/// - `list(...)[0]`
/// - `tuple(...)[0]`
/// - `list(i for i in ...)[0]`
/// - `[i for i in ...][0]`
/// - `list(...).pop(0)`
///
/// ## Why is this bad?
/// Calling `list(...)` will create a new list of the entire collection, which
/// can be very expensive for large collections. If you only need the first
/// element of the collection, you can use `next(...)` or `next(iter(...)` to
/// lazily fetch the first element.
/// Calling e.g. `list(...)` will create a new list of the entire collection,
/// which can be very expensive for large collections. If you only need the
/// first element of the collection, you can use `next(...)` or
/// `next(iter(...)` to lazily fetch the first element. The same is true for
/// the other constructs.
///
/// ## Example
/// ```python
Expand All @@ -33,14 +40,16 @@ use crate::fix::snippet::SourceCodeSnippet;
/// ```
///
/// ## Fix safety
/// This rule's fix is marked as unsafe, as migrating from `list(...)[0]` to
/// `next(iter(...))` can change the behavior of your program in two ways:
/// This rule's fix is marked as unsafe, as migrating from (e.g.) `list(...)[0]`
/// to `next(iter(...))` can change the behavior of your program in two ways:
///
/// 1. First, `list(...)` will eagerly evaluate the entire collection, while
/// `next(iter(...))` will only evaluate the first element. As such, any
/// side effects that occur during iteration will be delayed.
/// 2. Second, `list(...)[0]` will raise `IndexError` if the collection is
/// empty, while `next(iter(...))` will raise `StopIteration`.
/// 1. First, all above mentioned constructs will eagerly evaluate the entire
/// collection, while `next(iter(...))` will only evaluate the first
/// element. As such, any side effects that occur during iteration will be
/// delayed.
/// 2. Second, accessing members of a collection via square bracket notation
/// `[0]` of the `pop()` function will raise `IndexError` if the collection
/// is empty, while `next(iter(...))` will raise `StopIteration`.
///
/// ## References
/// - [Iterators and Iterables in Python: Run Efficient Iterations](https://realpython.com/python-iterators-iterables/#when-to-use-an-iterator-in-python)
Expand All @@ -67,18 +76,39 @@ impl AlwaysFixableViolation for UnnecessaryIterableAllocationForFirstElement {
/// RUF015
pub(crate) fn unnecessary_iterable_allocation_for_first_element(
checker: &mut Checker,
subscript: &ast::ExprSubscript,
expr: &Expr,
) {
let ast::ExprSubscript {
value,
slice,
range,
..
} = subscript;

if !is_head_slice(slice) {
return;
}
let value = match expr {
// Ex) `list(x)[0]`
Expr::Subscript(ast::ExprSubscript { value, slice, .. }) => {
if !is_zero(slice) {
return;
}
value
}
// Ex) `list(x).pop(0)`
Expr::Call(ast::ExprCall {
func, arguments, ..
}) => {
if !arguments.keywords.is_empty() {
return;
}
let [arg] = arguments.args.as_ref() else {
return;
};
if !is_zero(arg) {
return;
}
let Expr::Attribute(ast::ExprAttribute { value, attr, .. }) = func.as_ref() else {
return;
};
if !matches!(attr.as_str(), "pop") {
return;
}
value
}
_ => return,
};

let Some(target) = match_iteration_target(value, checker.semantic()) else {
return;
Expand All @@ -94,19 +124,19 @@ pub(crate) fn unnecessary_iterable_allocation_for_first_element(
UnnecessaryIterableAllocationForFirstElement {
iterable: SourceCodeSnippet::new(iterable.to_string()),
},
*range,
expr.range(),
);

diagnostic.set_fix(Fix::unsafe_edit(Edit::range_replacement(
format!("next({iterable})"),
*range,
expr.range(),
)));

checker.diagnostics.push(diagnostic);
}

/// Check that the slice [`Expr`] is a slice of the first element (e.g., `x[0]`).
fn is_head_slice(expr: &Expr) -> bool {
fn is_zero(expr: &Expr) -> bool {
matches!(
expr,
Expr::NumberLiteral(ast::ExprNumberLiteral {
Expand Down
Expand Up @@ -383,14 +383,16 @@ RUF015.py:57:1: RUF015 [*] Prefer `next(zip(x, y))` over single element slice
57 |+next(zip(x, y))
58 58 | [*zip(x, y)][0]
59 59 |
60 60 |
60 60 | # RUF015 (pop)

RUF015.py:58:1: RUF015 [*] Prefer `next(zip(x, y))` over single element slice
|
56 | # RUF015 (zip)
57 | list(zip(x, y))[0]
58 | [*zip(x, y)][0]
| ^^^^^^^^^^^^^^^ RUF015
59 |
60 | # RUF015 (pop)
|
= help: Replace with `next(zip(x, y))`

Expand All @@ -401,23 +403,82 @@ RUF015.py:58:1: RUF015 [*] Prefer `next(zip(x, y))` over single element slice
58 |-[*zip(x, y)][0]
58 |+next(zip(x, y))
59 59 |
60 60 |
61 61 | def test():
60 60 | # RUF015 (pop)
61 61 | list(x).pop(0)

RUF015.py:63:5: RUF015 [*] Prefer `next(iter(zip(x, y)))` over single element slice
RUF015.py:61:1: RUF015 [*] Prefer `next(iter(x))` over single element slice
|
61 | def test():
62 | zip = list # Overwrite the builtin zip
63 | list(zip(x, y))[0]
| ^^^^^^^^^^^^^^^^^^ RUF015
60 | # RUF015 (pop)
61 | list(x).pop(0)
| ^^^^^^^^^^^^^^ RUF015
62 | [i for i in x].pop(0)
63 | list(i for i in x).pop(0)
|
= help: Replace with `next(iter(zip(x, y)))`
= help: Replace with `next(iter(x))`

Unsafe fix
60 60 |
61 61 | def test():
62 62 | zip = list # Overwrite the builtin zip
63 |- list(zip(x, y))[0]
63 |+ next(iter(zip(x, y)))
58 58 | [*zip(x, y)][0]
59 59 |
60 60 | # RUF015 (pop)
61 |-list(x).pop(0)
61 |+next(iter(x))
62 62 | [i for i in x].pop(0)
63 63 | list(i for i in x).pop(0)
64 64 |

RUF015.py:62:1: RUF015 [*] Prefer `next(iter(x))` over single element slice
|
60 | # RUF015 (pop)
61 | list(x).pop(0)
62 | [i for i in x].pop(0)
| ^^^^^^^^^^^^^^^^^^^^^ RUF015
63 | list(i for i in x).pop(0)
|
= help: Replace with `next(iter(x))`

Unsafe fix
59 59 |
60 60 | # RUF015 (pop)
61 61 | list(x).pop(0)
62 |-[i for i in x].pop(0)
62 |+next(iter(x))
63 63 | list(i for i in x).pop(0)
64 64 |
65 65 | # OK

RUF015.py:63:1: RUF015 [*] Prefer `next(iter(x))` over single element slice
|
61 | list(x).pop(0)
62 | [i for i in x].pop(0)
63 | list(i for i in x).pop(0)
| ^^^^^^^^^^^^^^^^^^^^^^^^^ RUF015
64 |
65 | # OK
|
= help: Replace with `next(iter(x))`

Unsafe fix
60 60 | # RUF015 (pop)
61 61 | list(x).pop(0)
62 62 | [i for i in x].pop(0)
63 |-list(i for i in x).pop(0)
63 |+next(iter(x))
64 64 |
65 65 | # OK
66 66 | list(x).pop(1)

RUF015.py:73:5: RUF015 [*] Prefer `next(iter(zip(x, y)))` over single element slice
|
71 | def test():
72 | zip = list # Overwrite the builtin zip
73 | list(zip(x, y))[0]
| ^^^^^^^^^^^^^^^^^^ RUF015
|
= help: Replace with `next(iter(zip(x, y)))`

Unsafe fix
70 70 |
71 71 | def test():
72 72 | zip = list # Overwrite the builtin zip
73 |- list(zip(x, y))[0]
73 |+ next(iter(zip(x, y)))

0 comments on commit 0851d42

Please sign in to comment.