diff --git a/crates/ruff/resources/test/fixtures/ruff/RUF015.py b/crates/ruff/resources/test/fixtures/ruff/RUF015.py index e19602de8dce65..393694d89b3f17 100644 --- a/crates/ruff/resources/test/fixtures/ruff/RUF015.py +++ b/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] diff --git a/crates/ruff/src/rules/ruff/rules/invalid_index_type.rs b/crates/ruff/src/rules/ruff/rules/invalid_index_type.rs index b8e832a62eef6e..7ff8b762ac109a 100644 --- a/crates/ruff/src/rules/ruff/rules/invalid_index_type.rs +++ b/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, diff --git a/crates/ruff/src/rules/ruff/rules/unnecessary_iterable_allocation_for_first_element.rs b/crates/ruff/src/rules/ruff/rules/unnecessary_iterable_allocation_for_first_element.rs index 17fa0922e532fb..ba7ff351cee466 100644 --- a/crates/ruff/src/rules/ruff/rules/unnecessary_iterable_allocation_for_first_element.rs +++ b/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,6 +70,7 @@ 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}))]"), @@ -75,10 +78,21 @@ impl AlwaysAutofixableViolation for UnnecessaryIterableAllocationForFirstElement } } +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 { } } -/// 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 { 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) -> 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 { + 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, diff --git a/crates/ruff/src/rules/ruff/snapshots/ruff__rules__ruff__tests__RUF015_RUF015.py.snap b/crates/ruff/src/rules/ruff/snapshots/ruff__rules__ruff__tests__RUF015_RUF015.py.snap index 4eb011f0fe48ac..536a8c7d051d4d 100644 --- a/crates/ruff/src/rules/ruff/snapshots/ruff__rules__ruff__tests__RUF015_RUF015.py.snap +++ b/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 |+)) +