[ruff] Expand rule for list(iterable).pop(0) idiom (RUF015) (#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 https://github.com/astral-sh/ruff/issues/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.
This commit is contained in:
Robin Caloudis 2024-02-28 01:24:28 +01:00 committed by GitHub
parent 8dc22d5793
commit a1e8784207
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
4 changed files with 146 additions and 42 deletions

View file

@ -57,6 +57,16 @@ revision_heads_map_ast = [
list(zip(x, y))[0] list(zip(x, y))[0]
[*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(): def test():
zip = list # Overwrite the builtin zip zip = list # Overwrite the builtin zip

View file

@ -116,7 +116,7 @@ pub(crate) fn expression(expr: &Expr, checker: &mut Checker) {
flake8_simplify::rules::use_capital_environment_variables(checker, expr); flake8_simplify::rules::use_capital_environment_variables(checker, expr);
} }
if checker.enabled(Rule::UnnecessaryIterableAllocationForFirstElement) { 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) { if checker.enabled(Rule::InvalidIndexType) {
ruff::rules::invalid_index_type(checker, subscript); ruff::rules::invalid_index_type(checker, subscript);
@ -965,6 +965,9 @@ pub(crate) fn expression(expr: &Expr, checker: &mut Checker) {
if checker.enabled(Rule::DefaultFactoryKwarg) { if checker.enabled(Rule::DefaultFactoryKwarg) {
ruff::rules::default_factory_kwarg(checker, call); 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) => { Expr::Dict(dict) => {
if checker.any_enabled(&[ if checker.any_enabled(&[

View file

@ -11,14 +11,21 @@ use crate::checkers::ast::Checker;
use crate::fix::snippet::SourceCodeSnippet; use crate::fix::snippet::SourceCodeSnippet;
/// ## What it does /// ## What it does
/// Checks for uses of `list(...)[0]` that can be replaced with /// Checks the following constructs, all of which can be replaced by
/// `next(iter(...))`. /// `next(iter(...))`:
///
/// - `list(...)[0]`
/// - `tuple(...)[0]`
/// - `list(i for i in ...)[0]`
/// - `[i for i in ...][0]`
/// - `list(...).pop(0)`
/// ///
/// ## Why is this bad? /// ## Why is this bad?
/// Calling `list(...)` will create a new list of the entire collection, which /// Calling e.g. `list(...)` will create a new list of the entire collection,
/// can be very expensive for large collections. If you only need the first /// which can be very expensive for large collections. If you only need the
/// element of the collection, you can use `next(...)` or `next(iter(...)` to /// first element of the collection, you can use `next(...)` or
/// lazily fetch the first element. /// `next(iter(...)` to lazily fetch the first element. The same is true for
/// the other constructs.
/// ///
/// ## Example /// ## Example
/// ```python /// ```python
@ -33,14 +40,16 @@ use crate::fix::snippet::SourceCodeSnippet;
/// ``` /// ```
/// ///
/// ## Fix safety /// ## Fix safety
/// This rule's fix is marked as unsafe, as migrating from `list(...)[0]` to /// This rule's fix is marked as unsafe, as migrating from (e.g.) `list(...)[0]`
/// `next(iter(...))` can change the behavior of your program in two ways: /// to `next(iter(...))` can change the behavior of your program in two ways:
/// ///
/// 1. First, `list(...)` will eagerly evaluate the entire collection, while /// 1. First, all above mentioned constructs will eagerly evaluate the entire
/// `next(iter(...))` will only evaluate the first element. As such, any /// collection, while `next(iter(...))` will only evaluate the first
/// side effects that occur during iteration will be delayed. /// element. As such, any side effects that occur during iteration will be
/// 2. Second, `list(...)[0]` will raise `IndexError` if the collection is /// delayed.
/// empty, while `next(iter(...))` will raise `StopIteration`. /// 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 /// ## References
/// - [Iterators and Iterables in Python: Run Efficient Iterations](https://realpython.com/python-iterators-iterables/#when-to-use-an-iterator-in-python) /// - [Iterators and Iterables in Python: Run Efficient Iterations](https://realpython.com/python-iterators-iterables/#when-to-use-an-iterator-in-python)
@ -67,18 +76,39 @@ impl AlwaysFixableViolation for UnnecessaryIterableAllocationForFirstElement {
/// RUF015 /// RUF015
pub(crate) fn unnecessary_iterable_allocation_for_first_element( pub(crate) fn unnecessary_iterable_allocation_for_first_element(
checker: &mut Checker, checker: &mut Checker,
subscript: &ast::ExprSubscript, expr: &Expr,
) { ) {
let ast::ExprSubscript { let value = match expr {
value, // Ex) `list(x)[0]`
slice, Expr::Subscript(ast::ExprSubscript { value, slice, .. }) => {
range, if !is_zero(slice) {
.. return;
} = subscript; }
value
if !is_head_slice(slice) { }
return; // 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 { let Some(target) = match_iteration_target(value, checker.semantic()) else {
return; return;
@ -94,19 +124,19 @@ pub(crate) fn unnecessary_iterable_allocation_for_first_element(
UnnecessaryIterableAllocationForFirstElement { UnnecessaryIterableAllocationForFirstElement {
iterable: SourceCodeSnippet::new(iterable.to_string()), iterable: SourceCodeSnippet::new(iterable.to_string()),
}, },
*range, expr.range(),
); );
diagnostic.set_fix(Fix::unsafe_edit(Edit::range_replacement( diagnostic.set_fix(Fix::unsafe_edit(Edit::range_replacement(
format!("next({iterable})"), format!("next({iterable})"),
*range, expr.range(),
))); )));
checker.diagnostics.push(diagnostic); checker.diagnostics.push(diagnostic);
} }
/// Check that the slice [`Expr`] is a slice of the first element (e.g., `x[0]`). /// 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!( matches!(
expr, expr,
Expr::NumberLiteral(ast::ExprNumberLiteral { Expr::NumberLiteral(ast::ExprNumberLiteral {

View file

@ -383,7 +383,7 @@ RUF015.py:57:1: RUF015 [*] Prefer `next(zip(x, y))` over single element slice
57 |+next(zip(x, y)) 57 |+next(zip(x, y))
58 58 | [*zip(x, y)][0] 58 58 | [*zip(x, y)][0]
59 59 | 59 59 |
60 60 | 60 60 | # RUF015 (pop)
RUF015.py:58:1: RUF015 [*] Prefer `next(zip(x, y))` over single element slice RUF015.py:58:1: RUF015 [*] Prefer `next(zip(x, y))` over single element slice
| |
@ -391,6 +391,8 @@ RUF015.py:58:1: RUF015 [*] Prefer `next(zip(x, y))` over single element slice
57 | list(zip(x, y))[0] 57 | list(zip(x, y))[0]
58 | [*zip(x, y)][0] 58 | [*zip(x, y)][0]
| ^^^^^^^^^^^^^^^ RUF015 | ^^^^^^^^^^^^^^^ RUF015
59 |
60 | # RUF015 (pop)
| |
= help: Replace with `next(zip(x, y))` = help: Replace with `next(zip(x, y))`
@ -401,23 +403,82 @@ RUF015.py:58:1: RUF015 [*] Prefer `next(zip(x, y))` over single element slice
58 |-[*zip(x, y)][0] 58 |-[*zip(x, y)][0]
58 |+next(zip(x, y)) 58 |+next(zip(x, y))
59 59 | 59 59 |
60 60 | 60 60 | # RUF015 (pop)
61 61 | def test(): 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(): 60 | # RUF015 (pop)
62 | zip = list # Overwrite the builtin zip 61 | list(x).pop(0)
63 | list(zip(x, y))[0] | ^^^^^^^^^^^^^^ RUF015
62 | [i for i in x].pop(0)
63 | list(i for i in x).pop(0)
|
= help: Replace with `next(iter(x))`
Unsafe fix
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 | ^^^^^^^^^^^^^^^^^^ RUF015
| |
= help: Replace with `next(iter(zip(x, y)))` = help: Replace with `next(iter(zip(x, y)))`
Unsafe fix Unsafe fix
60 60 | 70 70 |
61 61 | def test(): 71 71 | def test():
62 62 | zip = list # Overwrite the builtin zip 72 72 | zip = list # Overwrite the builtin zip
63 |- list(zip(x, y))[0] 73 |- list(zip(x, y))[0]
63 |+ next(iter(zip(x, y))) 73 |+ next(iter(zip(x, y)))