Expand RUF015 to include all expression types

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

@@ -34,11 +34,23 @@
 [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]
+
+# RUF015 (multi-line)
+revision_heads_map_ast = [
+    a
+    for a in revision_heads_map_ast_obj.body
+    if isinstance(a, ast.Assign) and a.targets[0].id == "REVISION_HEADS_MAP"
+][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, TextSize};
+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
@@ -53,6 +54,7 @@ impl AlwaysAutofixableViolation for UnnecessaryIterableAllocationForFirstElement
             iterable,
             subscript_kind,
         } = self;
+        let iterable = Self::truncate(iterable);
         match subscript_kind {
             HeadSubscriptKind::Index => {
                 format!("Prefer `next(iter({iterable}))` over `list({iterable})[0]`")
@@ -68,17 +70,29 @@ impl AlwaysAutofixableViolation for UnnecessaryIterableAllocationForFirstElement
             iterable,
             subscript_kind,
         } = self;
+        let iterable = Self::truncate(iterable);
         match subscript_kind {
             HeadSubscriptKind::Index => format!("Replace with `next(iter({iterable}))`"),
             HeadSubscriptKind::Slice => format!("Replace with `[next(iter({iterable}))]"),
         }
     }
 }
 
+impl UnnecessaryIterableAllocationForFirstElement {
+    /// If the iterable is too long, or spans multiple lines, truncate it.
+    fn truncate(iterable: &str) -> &str {
+        if iterable.len() > 40 || iterable.contains(['\r', '\n']) {
+            "..."
+        } else {
+            iterable
+        }
+    }
+}
+
 /// RUF015
 pub(crate) fn unnecessary_iterable_allocation_for_first_element(
     checker: &mut Checker,
-    subscript: &ExprSubscript,
+    subscript: &ast::ExprSubscript,
 ) {
     let ast::ExprSubscript {
         value,
@@ -91,7 +105,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;
     };
 
@@ -164,9 +180,15 @@ 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> {
+/// Return the [`TextRange`] of the iterable from an expression, if the expression can be sliced
+/// into (i.e., is a list comprehension, or call to `list` or `tuple`).
+///
+/// For example, given `list(x)`, returns the range of `x`. Given `[x * x for x in y]`, returns the
+/// range of `x * x for x in y`.
+///
+/// As a special-case, given `[x for x in y]`, returns the range of `y` (rather than the
+/// redundant comprehension).
+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,43 +201,67 @@ 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 {
+            let [arg] = args.as_slice() else {
+                return None;
+            };
+
+            match arg {
+                Expr::GeneratorExp(ast::ExprGeneratorExp {
                     elt, generators, ..
-                })) => generator_iterable(elt, generators),
-                _ => None,
+                }) => match_simple_comprehension(elt, generators).or_else(|| Some(arg.range())),
+                Expr::ListComp(ast::ExprListComp {
+                    elt, generators, ..
+                }) => match_simple_comprehension(elt, generators).or_else(|| {
+                    Some(
+                        arg.range()
+                            // Remove the `[`
+                            .add_start(TextSize::from(1))
+                            // Remove the `]`
+                            .sub_end(TextSize::from(1)),
+                    )
+                }),
+                _ => Some(arg.range()),
             }
         }
         Expr::ListComp(ast::ExprListComp {
             elt, generators, ..
-        }) => generator_iterable(elt, generators),
+        }) => match_simple_comprehension(elt, generators).or_else(|| {
+            Some(
+                expr.range()
+                    // Remove the `[`
+                    .add_start(TextSize::from(1))
+                    // Remove the `]`
+                    .sub_end(TextSize::from(1)),
+            )
+        }),
         _ => 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> {
-    // 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() {
+/// Returns the [`Expr`] target for a comprehension, if the comprehension is "simple"
+/// (e.g., `x` for `[i for i in x]`).
+fn match_simple_comprehension(elt: &Expr, generators: &[Comprehension]) -> Option<TextRange> {
+    let [generator] = generators else {
         return None;
-    }
+    };
 
-    // 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 {
+    if generator.is_async {
         return None;
-    };
+    }
 
     // Ignore if there's an `if` statement in the comprehension, since it filters the list.
     if !generator.ifs.is_empty() {
         return None;
     }
 
-    let ast::ExprName { id, .. } = generator.iter.as_name_expr()?;
-    Some(id.as_str())
+    // Verify that the generator is, e.g. `i for i in x`, as opposed to `i for j in x`.
+    let elt = elt.as_name_expr()?;
+    let target = generator.target.as_name_expr()?;
+    if elt.id != target.id {
+        return None;
+    }
+
+    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,172 @@ 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 for j in y))` over `list(i + j for i in x for j in y)[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 for j in y))`
+
+ℹ 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 for j in y))
+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]
+50 50 | 
+
+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]
+50 50 | 
+51 51 | # RUF015 (multi-line)
+
+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
+50 | 
+51 | # RUF015 (multi-line)
+   |
+   = 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"]))
+50 50 | 
+51 51 | # RUF015 (multi-line)
+52 52 | revision_heads_map_ast = [
+
+RUF015.py:52:26: RUF015 [*] Prefer `next(iter(...))` over `list(...)[0]`
+   |
+51 |   # RUF015 (multi-line)
+52 |   revision_heads_map_ast = [
+   |  __________________________^
+53 | |     a
+54 | |     for a in revision_heads_map_ast_obj.body
+55 | |     if isinstance(a, ast.Assign) and a.targets[0].id == "REVISION_HEADS_MAP"
+56 | | ][0]
+   | |____^ RUF015
+   |
+   = help: Replace with `next(iter(...))`
+
+ℹ Suggested fix
+49 49 | list(x["y"])[0]
+50 50 | 
+51 51 | # RUF015 (multi-line)
+52    |-revision_heads_map_ast = [
+   52 |+revision_heads_map_ast = next(iter(
+53 53 |     a
+54 54 |     for a in revision_heads_map_ast_obj.body
+55 55 |     if isinstance(a, ast.Assign) and a.targets[0].id == "REVISION_HEADS_MAP"
+56    |-][0]
+   56 |+))
+