[ruff] Expand rule for list(iterable).pop(0) idiom (RUF015)

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

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(&[

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 {

crates/ruff_linter/src/rules/ruff/snapshots/ruff_linter__rules__ruff__tests__RUF015_RUF015.py.snap

@@ -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))`
 
@@ -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)))