Expand RUF015 to include all expression types

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

@@ -34,11 +34,16 @@
 [i for i in x][::2]
 [i for i in x][::]
 
-# OK (doesn't mirror the underlying list)
+# RUF015 (doesn't mirror the underlying list)
 [i + 1 for i in x][0]
 [i for i in x if i > 5][0]
 [(i, i + 1) for i in x][0]
 
-# OK (multiple generators)
+# RUF015 (multiple generators)
 y = range(10)
 [i + j for i in x for j in y][0]
+
+# RUF015
+list(range(10))[0]
+list(x.y)[0]
+list(x["y"])[0]

crates/ruff/src/rules/ruff/rules/invalid_index_type.rs

@@ -49,7 +49,7 @@ impl Violation for InvalidIndexType {
     }
 }
 
-/// RUF015
+/// RUF016
 pub(crate) fn invalid_index_type(checker: &mut Checker, expr: &ExprSubscript) {
     let ExprSubscript {
         value,

crates/ruff/src/rules/ruff/rules/unnecessary_iterable_allocation_for_first_element.rs

@@ -1,6 +1,7 @@
 use num_bigint::BigInt;
 use num_traits::{One, Zero};
-use rustpython_parser::ast::{self, Comprehension, Constant, Expr, ExprSubscript};
+use ruff_text_size::TextRange;
+use rustpython_parser::ast::{self, Comprehension, Constant, Expr, Ranged};
 
 use ruff_diagnostics::{AlwaysAutofixableViolation, Diagnostic, Edit, Fix};
 use ruff_macros::{derive_message_formats, violation};
@@ -30,12 +31,12 @@ use crate::registry::AsRule;
 ///
 /// ## Example
 /// ```python
-/// head = list(range(1000000000000))[0]
+/// head = list(x)[0]
 /// ```
 ///
 /// Use instead:
 /// ```python
-/// head = next(iter(range(1000000000000)))
+/// head = next(iter(x))
 /// ```
 ///
 /// ## References
@@ -78,7 +79,7 @@ impl AlwaysAutofixableViolation for UnnecessaryIterableAllocationForFirstElement
 /// RUF015
 pub(crate) fn unnecessary_iterable_allocation_for_first_element(
     checker: &mut Checker,
-    subscript: &ExprSubscript,
+    subscript: &ast::ExprSubscript,
 ) {
     let ast::ExprSubscript {
         value,
@@ -91,7 +92,9 @@ pub(crate) fn unnecessary_iterable_allocation_for_first_element(
         return;
     };
 
-    let Some(iterable) = iterable_name(value, checker.semantic()) else {
+    let Some(iterable) =
+        match_iterable(value, checker.semantic()).map(|range| checker.locator.slice(range))
+    else {
         return;
     };
 
@@ -166,7 +169,7 @@ fn classify_subscript(expr: &Expr) -> Option<HeadSubscriptKind> {
 
 /// Fetch the name of the iterable from an expression if the expression returns an unmodified list
 /// which can be sliced into.
-fn iterable_name<'a>(expr: &'a Expr, model: &SemanticModel) -> Option<&'a str> {
+fn match_iterable(expr: &Expr, model: &SemanticModel) -> Option<TextRange> {
     match expr {
         Expr::Call(ast::ExprCall { func, args, .. }) => {
             let ast::ExprName { id, .. } = func.as_name_expr()?;
@@ -179,24 +182,42 @@ fn iterable_name<'a>(expr: &'a Expr, model: &SemanticModel) -> Option<&'a str> {
                 return None;
             }
 
-            match args.first() {
-                Some(Expr::Name(ast::ExprName { id: arg_name, .. })) => Some(arg_name.as_str()),
-                Some(Expr::GeneratorExp(ast::ExprGeneratorExp {
-                    elt, generators, ..
-                })) => generator_iterable(elt, generators),
-                _ => None,
+            let [arg] = args.as_slice() else {
+                return None;
+            };
+
+            if let Expr::GeneratorExp(ast::ExprGeneratorExp {
+                elt, generators, ..
+            })
+            | Expr::ListComp(ast::ExprListComp {
+                elt, generators, ..
+            }) = arg
+            {
+                match_generator(elt, generators)
+            } else {
+                Some(arg.range())
             }
         }
         Expr::ListComp(ast::ExprListComp {
             elt, generators, ..
-        }) => generator_iterable(elt, generators),
+        }) => match_generator(elt, generators),
         _ => None,
     }
 }
 
 /// Given a comprehension, returns the name of the iterable over which it iterates, if it's
 /// a simple comprehension (e.g., `x` for `[i for i in x]`).
-fn generator_iterable<'a>(elt: &'a Expr, generators: &'a Vec<Comprehension>) -> Option<&'a str> {
+fn match_generator(elt: &Expr, generators: &[Comprehension]) -> Option<TextRange> {
+    match_simple_generator(elt, generators).or_else(|| {
+        let start = elt.start();
+        let end = generators.iter().map(Ranged::end).min()?;
+        Some(TextRange::new(start, end))
+    })
+}
+
+/// Given a comprehension, returns the name of the iterable over which it iterates, if it's
+/// a simple comprehension (e.g., `x` for `[i for i in x]`).
+fn match_simple_generator(elt: &Expr, generators: &[Comprehension]) -> Option<TextRange> {
     // 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 (e.g., `[i + 1 for i in x][0]`).
     if !elt.is_name_expr() {
@@ -205,7 +226,7 @@ fn generator_iterable<'a>(elt: &'a Expr, generators: &'a Vec<Comprehension>) ->
 
     // If there's more than 1 generator, we can't safely say that it fits the diagnostic conditions
     // (e.g., `[(i, j) for i in x for j in y][0]`).
-    let [generator] = generators.as_slice() else {
+    let [generator] = generators else {
         return None;
     };
 
@@ -214,8 +235,7 @@ fn generator_iterable<'a>(elt: &'a Expr, generators: &'a Vec<Comprehension>) ->
         return None;
     }
 
-    let ast::ExprName { id, .. } = generator.iter.as_name_expr()?;
-    Some(id.as_str())
+    Some(generator.iter.range())
 }
 
 /// If an expression is a constant integer, returns the value of that integer; otherwise,

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

@@ -335,4 +335,139 @@ RUF015.py:19:1: RUF015 [*] Prefer `[next(iter(x))]` over `list(x)[:1]`
 21 21 | # OK (not indexing (solely) the first element)
 22 22 | list(x)
 
+RUF015.py:38:1: RUF015 [*] Prefer `next(iter(i + 1 for i in x))` over `list(i + 1 for i in x)[0]`
+   |
+37 | # RUF015 (doesn't mirror the underlying list)
+38 | [i + 1 for i in x][0]
+   | ^^^^^^^^^^^^^^^^^^^^^ RUF015
+39 | [i for i in x if i > 5][0]
+40 | [(i, i + 1) for i in x][0]
+   |
+   = help: Replace with `next(iter(i + 1 for i in x))`
+
+ℹ Suggested fix
+35 35 | [i for i in x][::]
+36 36 | 
+37 37 | # RUF015 (doesn't mirror the underlying list)
+38    |-[i + 1 for i in x][0]
+   38 |+next(iter(i + 1 for i in x))
+39 39 | [i for i in x if i > 5][0]
+40 40 | [(i, i + 1) for i in x][0]
+41 41 | 
+
+RUF015.py:39:1: RUF015 [*] Prefer `next(iter(i for i in x if i > 5))` over `list(i for i in x if i > 5)[0]`
+   |
+37 | # RUF015 (doesn't mirror the underlying list)
+38 | [i + 1 for i in x][0]
+39 | [i for i in x if i > 5][0]
+   | ^^^^^^^^^^^^^^^^^^^^^^^^^^ RUF015
+40 | [(i, i + 1) for i in x][0]
+   |
+   = help: Replace with `next(iter(i for i in x if i > 5))`
+
+ℹ Suggested fix
+36 36 | 
+37 37 | # RUF015 (doesn't mirror the underlying list)
+38 38 | [i + 1 for i in x][0]
+39    |-[i for i in x if i > 5][0]
+   39 |+next(iter(i for i in x if i > 5))
+40 40 | [(i, i + 1) for i in x][0]
+41 41 | 
+42 42 | # RUF015 (multiple generators)
+
+RUF015.py:40:1: RUF015 [*] Prefer `next(iter((i, i + 1) for i in x))` over `list((i, i + 1) for i in x)[0]`
+   |
+38 | [i + 1 for i in x][0]
+39 | [i for i in x if i > 5][0]
+40 | [(i, i + 1) for i in x][0]
+   | ^^^^^^^^^^^^^^^^^^^^^^^^^^ RUF015
+41 | 
+42 | # RUF015 (multiple generators)
+   |
+   = help: Replace with `next(iter((i, i + 1) for i in x))`
+
+ℹ Suggested fix
+37 37 | # RUF015 (doesn't mirror the underlying list)
+38 38 | [i + 1 for i in x][0]
+39 39 | [i for i in x if i > 5][0]
+40    |-[(i, i + 1) for i in x][0]
+   40 |+next(iter((i, i + 1) for i in x))
+41 41 | 
+42 42 | # RUF015 (multiple generators)
+43 43 | y = range(10)
+
+RUF015.py:44:1: RUF015 [*] Prefer `next(iter(i + j for i in x))` over `list(i + j for i in x)[0]`
+   |
+42 | # RUF015 (multiple generators)
+43 | y = range(10)
+44 | [i + j for i in x for j in y][0]
+   | ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ RUF015
+45 | 
+46 | # RUF015
+   |
+   = help: Replace with `next(iter(i + j for i in x))`
+
+ℹ Suggested fix
+41 41 | 
+42 42 | # RUF015 (multiple generators)
+43 43 | y = range(10)
+44    |-[i + j for i in x for j in y][0]
+   44 |+next(iter(i + j for i in x))
+45 45 | 
+46 46 | # RUF015
+47 47 | list(range(10))[0]
+
+RUF015.py:47:1: RUF015 [*] Prefer `next(iter(range(10)))` over `list(range(10))[0]`
+   |
+46 | # RUF015
+47 | list(range(10))[0]
+   | ^^^^^^^^^^^^^^^^^^ RUF015
+48 | list(x.y)[0]
+49 | list(x["y"])[0]
+   |
+   = help: Replace with `next(iter(range(10)))`
+
+ℹ Suggested fix
+44 44 | [i + j for i in x for j in y][0]
+45 45 | 
+46 46 | # RUF015
+47    |-list(range(10))[0]
+   47 |+next(iter(range(10)))
+48 48 | list(x.y)[0]
+49 49 | list(x["y"])[0]
+
+RUF015.py:48:1: RUF015 [*] Prefer `next(iter(x.y))` over `list(x.y)[0]`
+   |
+46 | # RUF015
+47 | list(range(10))[0]
+48 | list(x.y)[0]
+   | ^^^^^^^^^^^^ RUF015
+49 | list(x["y"])[0]
+   |
+   = help: Replace with `next(iter(x.y))`
+
+ℹ Suggested fix
+45 45 | 
+46 46 | # RUF015
+47 47 | list(range(10))[0]
+48    |-list(x.y)[0]
+   48 |+next(iter(x.y))
+49 49 | list(x["y"])[0]
+
+RUF015.py:49:1: RUF015 [*] Prefer `next(iter(x["y"]))` over `list(x["y"])[0]`
+   |
+47 | list(range(10))[0]
+48 | list(x.y)[0]
+49 | list(x["y"])[0]
+   | ^^^^^^^^^^^^^^^ RUF015
+   |
+   = help: Replace with `next(iter(x["y"]))`
+
+ℹ Suggested fix
+46 46 | # RUF015
+47 47 | list(range(10))[0]
+48 48 | list(x.y)[0]
+49    |-list(x["y"])[0]
+   49 |+next(iter(x["y"]))
+