diff --git a/crates/ruff_linter/resources/test/fixtures/ruff/RUF015.py b/crates/ruff_linter/resources/test/fixtures/ruff/RUF015.py index 9ce388f6ceaba..b813297ce7605 100644 --- a/crates/ruff_linter/resources/test/fixtures/ruff/RUF015.py +++ b/crates/ruff_linter/resources/test/fixtures/ruff/RUF015.py @@ -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 diff --git a/crates/ruff_linter/src/checkers/ast/analyze/expression.rs b/crates/ruff_linter/src/checkers/ast/analyze/expression.rs index 1e822a4775561..14385ce2bc00e 100644 --- a/crates/ruff_linter/src/checkers/ast/analyze/expression.rs +++ b/crates/ruff_linter/src/checkers/ast/analyze/expression.rs @@ -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); @@ -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(&[ diff --git a/crates/ruff_linter/src/rules/ruff/rules/unnecessary_iterable_allocation_for_first_element.rs b/crates/ruff_linter/src/rules/ruff/rules/unnecessary_iterable_allocation_for_first_element.rs index 9cac96d080794..6d4d0643b49a5 100644 --- a/crates/ruff_linter/src/rules/ruff/rules/unnecessary_iterable_allocation_for_first_element.rs +++ b/crates/ruff_linter/src/rules/ruff/rules/unnecessary_iterable_allocation_for_first_element.rs @@ -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 @@ -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) @@ -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; @@ -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 { diff --git a/crates/ruff_linter/src/rules/ruff/snapshots/ruff_linter__rules__ruff__tests__RUF015_RUF015.py.snap b/crates/ruff_linter/src/rules/ruff/snapshots/ruff_linter__rules__ruff__tests__RUF015_RUF015.py.snap index 016df57bdba9f..7500f39c81f82 100644 --- a/crates/ruff_linter/src/rules/ruff/snapshots/ruff_linter__rules__ruff__tests__RUF015_RUF015.py.snap +++ b/crates/ruff_linter/src/rules/ruff/snapshots/ruff_linter__rules__ruff__tests__RUF015_RUF015.py.snap @@ -383,7 +383,7 @@ 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 | @@ -391,6 +391,8 @@ RUF015.py:58:1: RUF015 [*] Prefer `next(zip(x, y))` over single element slice 57 | list(zip(x, y))[0] 58 | [*zip(x, y)][0] | ^^^^^^^^^^^^^^^ RUF015 +59 | +60 | # RUF015 (pop) | = 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 |+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)))