Fix range

astral-sh · Feb 28, 2024 · 2003ed7 · 2003ed7
1 parent 2ce293a
commit 2003ed7
Show file tree

Hide file tree

Showing 3 changed files with 63 additions and 24 deletions.
diff --git a/crates/ruff_linter/resources/test/fixtures/ruff/RUF015.py b/crates/ruff_linter/resources/test/fixtures/ruff/RUF015.py
@@ -69,3 +69,7 @@
 def test():
     zip = list  # Overwrite the builtin zip
     list(zip(x, y))[0]
+
+# RUF015 (pop)
+[i for i in x].pop(0)
+list(i for i in x).pop(0)
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
@@ -40,7 +40,7 @@ use crate::fix::snippet::SourceCodeSnippet;
 /// ```
 ///
 /// ## Fix safety
-/// This rule's fix is marked as unsafe, as migrating from e.g. `list(...)[0]`
+/// 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, all above mentioned constructs will eagerly evaluate the entire
@@ -76,39 +76,36 @@ impl AlwaysFixableViolation for UnnecessaryIterableAllocationForFirstElement {
 /// RUF015
 pub(crate) fn unnecessary_iterable_allocation_for_first_element(
     checker: &mut Checker,
-    expr: &ast::Expr,
+    expr: &Expr,
 ) {
-    let (value, range) = match expr {
-        ast::Expr::Subscript(ast::ExprSubscript {
-            value,
-            slice,
-            range,
-            ..
-        }) => {
-            if !is_head_slice(slice) {
+    let value = match expr {
+        // Ex) `list(x)[0]`
+        Expr::Subscript(ast::ExprSubscript { value, slice, .. }) => {
+            if !is_zero(slice) {
                 return;
             }
-            (value, range)
+            value
         }
-        ast::Expr::Call(ast::ExprCall {
+        // Ex) `list(x).pop(0)`
+        Expr::Call(ast::ExprCall {
             func, arguments, ..
         }) => {
-            let Some(arg) = arguments.args.first() else {
+            if !arguments.keywords.is_empty() {
+                return;
+            }
+            let [arg] = arguments.args.as_ref() else {
                 return;
             };
-            if !is_head_slice(arg) {
+            if !is_zero(arg) {
                 return;
             }
-            let ast::Expr::Attribute(ast::ExprAttribute {
-                range, value, attr, ..
-            }) = func.as_ref()
-            else {
+            let Expr::Attribute(ast::ExprAttribute { value, attr, .. }) = func.as_ref() else {
                 return;
             };
             if !matches!(attr.as_str(), "pop") {
                 return;
             }
-            (value, range)
+            value
         }
         _ => return,
     };
@@ -127,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/...ff_linter/src/rules/ruff/snapshots/ruff_linter__rules__ruff__tests__RUF015_RUF015.py.snap b/...ff_linter/src/rules/ruff/snapshots/ruff_linter__rules__ruff__tests__RUF015_RUF015.py.snap
@@ -410,7 +410,7 @@ RUF015.py:61:1: RUF015 [*] Prefer `next(iter(x))` over single element slice
    |
 60 | # RUF015 (pop)
 61 | list(x).pop(0)
-   | ^^^^^^^^^^^ RUF015
+   | ^^^^^^^^^^^^^^ RUF015
 62 | 
 63 | # OK
    |
@@ -421,7 +421,7 @@ RUF015.py:61:1: RUF015 [*] Prefer `next(iter(x))` over single element slice
 59 59 | 
 60 60 | # RUF015 (pop)
 61    |-list(x).pop(0)
-   61 |+next(iter(x))(0)
+   61 |+next(iter(x))
 62 62 | 
 63 63 | # OK
 64 64 | list(x).pop(1)
@@ -432,6 +432,8 @@ RUF015.py:71:5: RUF015 [*] Prefer `next(iter(zip(x, y)))` over single element sl
 70 |     zip = list  # Overwrite the builtin zip
 71 |     list(zip(x, y))[0]
    |     ^^^^^^^^^^^^^^^^^^ RUF015
+72 | 
+73 | # RUF015 (pop)
    |
    = help: Replace with `next(iter(zip(x, y)))`
 
@@ -441,3 +443,39 @@ RUF015.py:71:5: RUF015 [*] Prefer `next(iter(zip(x, y)))` over single element sl
 70 70 |     zip = list  # Overwrite the builtin zip
 71    |-    list(zip(x, y))[0]
    71 |+    next(iter(zip(x, y)))
+72 72 | 
+73 73 | # RUF015 (pop)
+74 74 | [i for i in x].pop(0)
+
+RUF015.py:74:1: RUF015 [*] Prefer `next(iter(x))` over single element slice
+   |
+73 | # RUF015 (pop)
+74 | [i for i in x].pop(0)
+   | ^^^^^^^^^^^^^^^^^^^^^ RUF015
+75 | list(i for i in x).pop(0)
+   |
+   = help: Replace with `next(iter(x))`
+
+ℹ Unsafe fix
+71 71 |     list(zip(x, y))[0]
+72 72 | 
+73 73 | # RUF015 (pop)
+74    |-[i for i in x].pop(0)
+   74 |+next(iter(x))
+75 75 | list(i for i in x).pop(0)
+
+RUF015.py:75:1: RUF015 [*] Prefer `next(iter(x))` over single element slice
+   |
+73 | # RUF015 (pop)
+74 | [i for i in x].pop(0)
+75 | list(i for i in x).pop(0)
+   | ^^^^^^^^^^^^^^^^^^^^^^^^^ RUF015
+   |
+   = help: Replace with `next(iter(x))`
+
+ℹ Unsafe fix
+72 72 | 
+73 73 | # RUF015 (pop)
+74 74 | [i for i in x].pop(0)
+75    |-list(i for i in x).pop(0)
+   75 |+next(iter(x))