Handle non-Name elts in Expr::ListComp

crates/ruff/resources/test/fixtures/ruff/RUF015.py

@@ -3,17 +3,31 @@
 # RUF015
 list(x)[0]
 list(x)[:1]
-
-# RUF015
+list(x)[:1:1]
+list(x)[:1:2]
 [i for i in x][0]
 [i for i in x][:1]
+[i for i in x][:1:1]
+[i for i in x][:1:2]
 
-y = range(10)
-# Fine - multiple generators
-[i + j for i in x for j in y][0]
-
-# Fine - not indexing at 0
+# Fine - not indexing (solely) the first element
+list(x)
 list(x)[1]
+list(x)[-1]
 list(x)[1:]
+list(x)[:3:2]
+list(x)[::2]
+[i for i in x]
 [i for i in x][1]
+[i for i in x][-1]
 [i for i in x][1:]
+[i for i in x][:3:2]
+[i for i in x][::2]
+
+# Fine - modifies the first element
+[i + 1 for i in x][0]
+
+y = range(10)
+# Fine - multiple generators
+[i + j for i in x for j in y][0]
+

crates/ruff/src/rules/ruff/mod.rs

@@ -30,6 +30,7 @@ mod tests {
     #[test_case(Rule::MutableDataclassDefault, Path::new("RUF008.py"))]
     #[test_case(Rule::PairwiseOverZipped, Path::new("RUF007.py"))]
     #[test_case(Rule::StaticKeyDictComprehension, Path::new("RUF011.py"))]
+    #[test_case(Rule::UnnecessaryListAllocationForFirstElement, Path::new("RUF015.py"))]
     #[cfg_attr(
         feature = "unreachable-code",
         test_case(Rule::UnreachableCode, Path::new("RUF014.py"))

crates/ruff/src/rules/ruff/rules/unnecessary_list_allocation_for_first_element.rs

@@ -90,12 +90,12 @@ pub(crate) fn unnecessary_list_allocation_for_first_element(
     checker.diagnostics.push(diagnostic);
 }
 
-/// Fetch the name of the iterable from a list expression.
-fn get_iterable_name<'a>(checker: &mut Checker, expr: &'a Expr) -> Option<&'a String> {
-    // Decompose.
-    let name = match expr {
+/// Fetch the name of the iterable from a list expression if the expression returns an unmodified list
+/// which can be sliced into.
+fn get_iterable_name<'a>(checker: &mut Checker, expr: &'a Expr) -> Option<&'a str> {
+    match expr {
         Expr::Call(ast::ExprCall { func, args, .. }) => {
-            let Expr::Name(ast::ExprName { id, .. }) = func.as_ref() else {
+            let Some(id) = get_name_id(func.as_ref()) else {
                 return None;
             };
             if !(id == "list" && checker.semantic().is_builtin("list")) {
@@ -106,46 +106,69 @@ fn get_iterable_name<'a>(checker: &mut Checker, expr: &'a Expr) -> Option<&'a St
                 return None;
             };
 
-            Some(arg_name)
+            Some(arg_name.as_str())
         }
-        Expr::ListComp(ast::ExprListComp { generators, .. }) => {
+        Expr::ListComp(ast::ExprListComp {
+            elt, generators, ..
+        }) => {
+            // If the `elt` field is anything other than a [`Expr::Name`], we can't be sure that it
+            // doesn't modify the elements of the underlying iterator - for example, `[i + 1 for i in x][0]`.
+            if !matches!(elt.as_ref(), Expr::Name(ast::ExprName { .. })) {
+                return None;
+            }
+
             // If there's more than 1 generator, we can't safely say that it fits the diagnostic conditions -
-            // for example, `[i + j for i in x for j in y][0]`
+            // for example, `[(i, j) for i in x for j in y][0]`.
             if generators.len() != 1 {
                 return None;
             }
 
             let generator = &generators[0];
-            let Expr::Name(ast::ExprName { id: arg_name, .. }) = &generator.iter else {
+            let Some(arg_name) = get_name_id(&generator.iter) else {
                 return None;
             };
-
             Some(arg_name)
         }
         _ => None,
-    };
-
-    name
+    }
 }
 
+/// Check that the slice [`Expr`] is functionally equivalent to slicing into the first element.
 fn indexes_first_element(expr: &Expr) -> bool {
     match expr {
-        Expr::Constant(ast::ExprConstant { .. }) => get_index_value(expr) == Some(0i64),
-        Expr::Slice(ast::ExprSlice { lower, upper, .. }) => {
-            let lower_index = lower.as_ref().and_then(|l| get_index_value(&l));
-            let upper_index = upper.as_ref().and_then(|u| get_index_value(&u));
-
-            if lower_index.is_none() || lower_index == Some(0i64) {
-                return upper_index == Some(1i64);
-            } else {
-                return false;
+        Expr::Constant(ast::ExprConstant { .. }) => {
+            get_effective_index(expr).unwrap_or(0i64) == 0i64
+        }
+        Expr::Slice(ast::ExprSlice {
+            step: step_value,
+            lower: lower_index,
+            upper: upper_index,
+            ..
+        }) => {
+            let lower = lower_index
+                .as_ref()
+                .and_then(|l| get_effective_index(l))
+                .unwrap_or(0i64);
+            let upper = upper_index
+                .as_ref()
+                .and_then(|u| get_effective_index(u))
+                .unwrap_or(i64::MAX);
+            let step = step_value
+                .as_ref()
+                .and_then(|s| get_effective_index(s))
+                .unwrap_or(1i64);
+
+            if lower == 0i64 && upper <= step {
+                return true;
             }
+
+            false
         }
         _ => false,
     }
 }
 
-fn get_index_value(expr: &Expr) -> Option<i64> {
+fn get_effective_index(expr: &Expr) -> Option<i64> {
     match expr {
         Expr::Constant(ast::ExprConstant {
             value: Constant::Int(value),
@@ -154,3 +177,10 @@ fn get_index_value(expr: &Expr) -> Option<i64> {
         _ => None,
     }
 }
+
+fn get_name_id(expr: &Expr) -> Option<&str> {
+    match expr {
+        Expr::Name(ast::ExprName { id, .. }) => Some(id),
+        _ => None,
+    }
+}

crates/ruff/src/rules/ruff/snapshots/ruff__rules__ruff__tests__RUF015_RUF015.py.snap

@@ -0,0 +1,170 @@
+---
+source: crates/ruff/src/rules/ruff/mod.rs
+---
+RUF015.py:4:1: RUF015 [*] Prefer `next(iter(x))` over `list(x)[0]` or equivalent list comprehension
+  |
+3 | # RUF015
+4 | list(x)[0]
+  | ^^^^^^^^^^ RUF015
+5 | list(x)[:1]
+6 | list(x)[:1:1]
+  |
+  = help: Convert `list(x)[0]` or equivalent list comprehension call to `next(iter(x))`
+
+ℹ Suggested fix
+1 1 | x = range(10)
+2 2 | 
+3 3 | # RUF015
+4   |-list(x)[0]
+  4 |+next(iter(x))
+5 5 | list(x)[:1]
+6 6 | list(x)[:1:1]
+7 7 | list(x)[:1:2]
+
+RUF015.py:5:1: RUF015 [*] Prefer `next(iter(x))` over `list(x)[0]` or equivalent list comprehension
+  |
+3 | # RUF015
+4 | list(x)[0]
+5 | list(x)[:1]
+  | ^^^^^^^^^^^ RUF015
+6 | list(x)[:1:1]
+7 | list(x)[:1:2]
+  |
+  = help: Convert `list(x)[0]` or equivalent list comprehension call to `next(iter(x))`
+
+ℹ Suggested fix
+2 2 | 
+3 3 | # RUF015
+4 4 | list(x)[0]
+5   |-list(x)[:1]
+  5 |+next(iter(x))
+6 6 | list(x)[:1:1]
+7 7 | list(x)[:1:2]
+8 8 | [i for i in x][0]
+
+RUF015.py:6:1: RUF015 [*] Prefer `next(iter(x))` over `list(x)[0]` or equivalent list comprehension
+  |
+4 | list(x)[0]
+5 | list(x)[:1]
+6 | list(x)[:1:1]
+  | ^^^^^^^^^^^^^ RUF015
+7 | list(x)[:1:2]
+8 | [i for i in x][0]
+  |
+  = help: Convert `list(x)[0]` or equivalent list comprehension call to `next(iter(x))`
+
+ℹ Suggested fix
+3 3 | # RUF015
+4 4 | list(x)[0]
+5 5 | list(x)[:1]
+6   |-list(x)[:1:1]
+  6 |+next(iter(x))
+7 7 | list(x)[:1:2]
+8 8 | [i for i in x][0]
+9 9 | [i for i in x][:1]
+
+RUF015.py:7:1: RUF015 [*] Prefer `next(iter(x))` over `list(x)[0]` or equivalent list comprehension
+  |
+5 | list(x)[:1]
+6 | list(x)[:1:1]
+7 | list(x)[:1:2]
+  | ^^^^^^^^^^^^^ RUF015
+8 | [i for i in x][0]
+9 | [i for i in x][:1]
+  |
+  = help: Convert `list(x)[0]` or equivalent list comprehension call to `next(iter(x))`
+
+ℹ Suggested fix
+4 4 | list(x)[0]
+5 5 | list(x)[:1]
+6 6 | list(x)[:1:1]
+7   |-list(x)[:1:2]
+  7 |+next(iter(x))
+8 8 | [i for i in x][0]
+9 9 | [i for i in x][:1]
+10 10 | [i for i in x][:1:1]
+
+RUF015.py:8:1: RUF015 [*] Prefer `next(iter(x))` over `list(x)[0]` or equivalent list comprehension
+   |
+ 6 | list(x)[:1:1]
+ 7 | list(x)[:1:2]
+ 8 | [i for i in x][0]
+   | ^^^^^^^^^^^^^^^^^ RUF015
+ 9 | [i for i in x][:1]
+10 | [i for i in x][:1:1]
+   |
+   = help: Convert `list(x)[0]` or equivalent list comprehension call to `next(iter(x))`
+
+ℹ Suggested fix
+5 5 | list(x)[:1]
+6 6 | list(x)[:1:1]
+7 7 | list(x)[:1:2]
+8   |-[i for i in x][0]
+  8 |+next(iter(x))
+9 9 | [i for i in x][:1]
+10 10 | [i for i in x][:1:1]
+11 11 | [i for i in x][:1:2]
+
+RUF015.py:9:1: RUF015 [*] Prefer `next(iter(x))` over `list(x)[0]` or equivalent list comprehension
+   |
+ 7 | list(x)[:1:2]
+ 8 | [i for i in x][0]
+ 9 | [i for i in x][:1]
+   | ^^^^^^^^^^^^^^^^^^ RUF015
+10 | [i for i in x][:1:1]
+11 | [i for i in x][:1:2]
+   |
+   = help: Convert `list(x)[0]` or equivalent list comprehension call to `next(iter(x))`
+
+ℹ Suggested fix
+6  6  | list(x)[:1:1]
+7  7  | list(x)[:1:2]
+8  8  | [i for i in x][0]
+9     |-[i for i in x][:1]
+   9  |+next(iter(x))
+10 10 | [i for i in x][:1:1]
+11 11 | [i for i in x][:1:2]
+12 12 | 
+
+RUF015.py:10:1: RUF015 [*] Prefer `next(iter(x))` over `list(x)[0]` or equivalent list comprehension
+   |
+ 8 | [i for i in x][0]
+ 9 | [i for i in x][:1]
+10 | [i for i in x][:1:1]
+   | ^^^^^^^^^^^^^^^^^^^^ RUF015
+11 | [i for i in x][:1:2]
+   |
+   = help: Convert `list(x)[0]` or equivalent list comprehension call to `next(iter(x))`
+
+ℹ Suggested fix
+7  7  | list(x)[:1:2]
+8  8  | [i for i in x][0]
+9  9  | [i for i in x][:1]
+10    |-[i for i in x][:1:1]
+   10 |+next(iter(x))
+11 11 | [i for i in x][:1:2]
+12 12 | 
+13 13 | # Fine - not indexing (solely) the first element
+
+RUF015.py:11:1: RUF015 [*] Prefer `next(iter(x))` over `list(x)[0]` or equivalent list comprehension
+   |
+ 9 | [i for i in x][:1]
+10 | [i for i in x][:1:1]
+11 | [i for i in x][:1:2]
+   | ^^^^^^^^^^^^^^^^^^^^ RUF015
+12 | 
+13 | # Fine - not indexing (solely) the first element
+   |
+   = help: Convert `list(x)[0]` or equivalent list comprehension call to `next(iter(x))`
+
+ℹ Suggested fix
+8  8  | [i for i in x][0]
+9  9  | [i for i in x][:1]
+10 10 | [i for i in x][:1:1]
+11    |-[i for i in x][:1:2]
+   11 |+next(iter(x))
+12 12 | 
+13 13 | # Fine - not indexing (solely) the first element
+14 14 | list(x)
+
+