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 ba7ff351cee466..138bab7879bf80 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,3 +1,5 @@ +use std::borrow::Cow; + use num_bigint::BigInt; use num_traits::{One, Zero}; use ruff_text_size::{TextRange, TextSize}; @@ -17,8 +19,8 @@ use crate::registry::AsRule; /// ## Why is this bad? /// Calling `list(...)` will create a new list of the entire collection, which /// can be very expensive for large collections. If you only need the first -/// element of the collection, you can use `next(iter(...))` to lazily fetch -/// the first element without creating a new list. +/// element of the collection, you can use `next(...)` or `next(iter(...)` to +/// lazily fetch the first element. /// /// Note that migrating from `list(...)[0]` to `next(iter(...))` can change /// the behavior of your program in two ways: @@ -27,16 +29,18 @@ use crate::registry::AsRule; /// `next(iter(...))` will only evaluate the first element. As such, any /// side effects that occur during iteration will be delayed. /// 2. Second, `list(...)[0]` will raise `IndexError` if the collection is -/// empty, while `next(iter(...))` will raise `StopIteration`. +/// empty, while `next(iter(...))` will raise `StopIteration`. /// /// ## Example /// ```python /// head = list(x)[0] +/// head = [x * x for x in range(10)][0] /// ``` /// /// Use instead: /// ```python /// head = next(iter(x)) +/// head = next(x * x for x in range(10)) /// ``` /// /// ## References @@ -57,10 +61,10 @@ impl AlwaysAutofixableViolation for UnnecessaryIterableAllocationForFirstElement let iterable = Self::truncate(iterable); match subscript_kind { HeadSubscriptKind::Index => { - format!("Prefer `next(iter({iterable}))` over `list({iterable})[0]`") + format!("Prefer `next({iterable})` over single element slice") } HeadSubscriptKind::Slice => { - format!("Prefer `[next(iter({iterable}))]` over `list({iterable})[:1]`") + format!("Prefer `[next({iterable})]` over single element slice") } } } @@ -72,8 +76,8 @@ impl AlwaysAutofixableViolation for UnnecessaryIterableAllocationForFirstElement } = 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}))]"), + HeadSubscriptKind::Index => format!("Replace with `next({iterable})`"), + HeadSubscriptKind::Slice => format!("Replace with `[next({iterable})]"), } } } @@ -105,11 +109,15 @@ pub(crate) fn unnecessary_iterable_allocation_for_first_element( return; }; - let Some(iterable) = - match_iterable(value, checker.semantic()).map(|range| checker.locator.slice(range)) - else { + let Some(target) = match_iteration_target(value, checker.semantic()) else { return; }; + let iterable = checker.locator.slice(target.range); + let iterable = if target.iterable { + Cow::Borrowed(iterable) + } else { + Cow::Owned(format!("iter({iterable})")) + }; let mut diagnostic = Diagnostic::new( UnnecessaryIterableAllocationForFirstElement { @@ -121,8 +129,8 @@ pub(crate) fn unnecessary_iterable_allocation_for_first_element( if checker.patch(diagnostic.kind.rule()) { let replacement = match subscript_kind { - HeadSubscriptKind::Index => format!("next(iter({iterable}))"), - HeadSubscriptKind::Slice => format!("[next(iter({iterable}))]"), + HeadSubscriptKind::Index => format!("next({iterable})"), + HeadSubscriptKind::Slice => format!("[next({iterable})]"), }; diagnostic.set_fix(Fix::suggested(Edit::range_replacement(replacement, *range))); } @@ -180,15 +188,24 @@ fn classify_subscript(expr: &Expr) -> Option { } } -/// 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`). +#[derive(Debug)] +struct IterationTarget { + /// The [`TextRange`] of the target. + range: TextRange, + /// Whether the target is an iterable (e.g., a generator). If not, the target must be wrapped + /// in `iter(...)` prior to calling `next(...)`. + iterable: bool, +} + +/// Return the [`IterationTarget`] of 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 { +fn match_iteration_target(expr: &Expr, model: &SemanticModel) -> Option { match expr { Expr::Call(ast::ExprCall { func, args, .. }) => { let ast::ExprName { id, .. } = func.as_name_expr()?; @@ -197,43 +214,70 @@ fn match_iterable(expr: &Expr, model: &SemanticModel) -> Option { return None; } - if !model.is_builtin(id.as_str()) { - return None; - } - let [arg] = args.as_slice() else { return None; }; + if !model.is_builtin(id.as_str()) { + return None; + } + match arg { Expr::GeneratorExp(ast::ExprGeneratorExp { elt, generators, .. - }) => match_simple_comprehension(elt, generators).or_else(|| Some(arg.range())), + }) => match_simple_comprehension(elt, generators) + .map(|range| IterationTarget { + range, + iterable: false, + }) + .or_else(|| { + Some(IterationTarget { + range: arg.range(), + iterable: true, + }) + }), 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)), - ) + }) => match_simple_comprehension(elt, generators) + .map(|range| IterationTarget { + range, + iterable: false, + }) + .or_else(|| { + Some(IterationTarget { + range: arg + .range() + // Remove the `[` + .add_start(TextSize::from(1)) + // Remove the `]` + .sub_end(TextSize::from(1)), + iterable: true, + }) + }), + _ => Some(IterationTarget { + range: arg.range(), + iterable: false, }), - _ => Some(arg.range()), } } Expr::ListComp(ast::ExprListComp { 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)), - ) - }), + }) => match_simple_comprehension(elt, generators) + .map(|range| IterationTarget { + range, + iterable: false, + }) + .or_else(|| { + Some(IterationTarget { + range: expr + .range() + // Remove the `[` + .add_start(TextSize::from(1)) + // Remove the `]` + .sub_end(TextSize::from(1)), + iterable: true, + }) + }), _ => None, } } 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 536a8c7d051d4d..c55e737c6d0b34 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 @@ -1,7 +1,7 @@ --- source: crates/ruff/src/rules/ruff/mod.rs --- -RUF015.py:4:1: RUF015 [*] Prefer `next(iter(x))` over `list(x)[0]` +RUF015.py:4:1: RUF015 [*] Prefer `next(iter(x))` over single element slice | 3 | # RUF015 4 | list(x)[0] @@ -21,7 +21,7 @@ RUF015.py:4:1: RUF015 [*] Prefer `next(iter(x))` over `list(x)[0]` 6 6 | list(x)[:1:1] 7 7 | list(x)[:1:2] -RUF015.py:5:1: RUF015 [*] Prefer `[next(iter(x))]` over `list(x)[:1]` +RUF015.py:5:1: RUF015 [*] Prefer `[next(iter(x))]` over single element slice | 3 | # RUF015 4 | list(x)[0] @@ -42,7 +42,7 @@ RUF015.py:5:1: RUF015 [*] Prefer `[next(iter(x))]` over `list(x)[:1]` 7 7 | list(x)[:1:2] 8 8 | tuple(x)[0] -RUF015.py:6:1: RUF015 [*] Prefer `[next(iter(x))]` over `list(x)[:1]` +RUF015.py:6:1: RUF015 [*] Prefer `[next(iter(x))]` over single element slice | 4 | list(x)[0] 5 | list(x)[:1] @@ -63,7 +63,7 @@ RUF015.py:6:1: RUF015 [*] Prefer `[next(iter(x))]` over `list(x)[:1]` 8 8 | tuple(x)[0] 9 9 | tuple(x)[:1] -RUF015.py:7:1: RUF015 [*] Prefer `[next(iter(x))]` over `list(x)[:1]` +RUF015.py:7:1: RUF015 [*] Prefer `[next(iter(x))]` over single element slice | 5 | list(x)[:1] 6 | list(x)[:1:1] @@ -84,7 +84,7 @@ RUF015.py:7:1: RUF015 [*] Prefer `[next(iter(x))]` over `list(x)[:1]` 9 9 | tuple(x)[:1] 10 10 | tuple(x)[:1:1] -RUF015.py:8:1: RUF015 [*] Prefer `next(iter(x))` over `list(x)[0]` +RUF015.py:8:1: RUF015 [*] Prefer `next(iter(x))` over single element slice | 6 | list(x)[:1:1] 7 | list(x)[:1:2] @@ -105,7 +105,7 @@ RUF015.py:8:1: RUF015 [*] Prefer `next(iter(x))` over `list(x)[0]` 10 10 | tuple(x)[:1:1] 11 11 | tuple(x)[:1:2] -RUF015.py:9:1: RUF015 [*] Prefer `[next(iter(x))]` over `list(x)[:1]` +RUF015.py:9:1: RUF015 [*] Prefer `[next(iter(x))]` over single element slice | 7 | list(x)[:1:2] 8 | tuple(x)[0] @@ -126,7 +126,7 @@ RUF015.py:9:1: RUF015 [*] Prefer `[next(iter(x))]` over `list(x)[:1]` 11 11 | tuple(x)[:1:2] 12 12 | list(i for i in x)[0] -RUF015.py:10:1: RUF015 [*] Prefer `[next(iter(x))]` over `list(x)[:1]` +RUF015.py:10:1: RUF015 [*] Prefer `[next(iter(x))]` over single element slice | 8 | tuple(x)[0] 9 | tuple(x)[:1] @@ -147,7 +147,7 @@ RUF015.py:10:1: RUF015 [*] Prefer `[next(iter(x))]` over `list(x)[:1]` 12 12 | list(i for i in x)[0] 13 13 | list(i for i in x)[:1] -RUF015.py:11:1: RUF015 [*] Prefer `[next(iter(x))]` over `list(x)[:1]` +RUF015.py:11:1: RUF015 [*] Prefer `[next(iter(x))]` over single element slice | 9 | tuple(x)[:1] 10 | tuple(x)[:1:1] @@ -168,7 +168,7 @@ RUF015.py:11:1: RUF015 [*] Prefer `[next(iter(x))]` over `list(x)[:1]` 13 13 | list(i for i in x)[:1] 14 14 | list(i for i in x)[:1:1] -RUF015.py:12:1: RUF015 [*] Prefer `next(iter(x))` over `list(x)[0]` +RUF015.py:12:1: RUF015 [*] Prefer `next(iter(x))` over single element slice | 10 | tuple(x)[:1:1] 11 | tuple(x)[:1:2] @@ -189,7 +189,7 @@ RUF015.py:12:1: RUF015 [*] Prefer `next(iter(x))` over `list(x)[0]` 14 14 | list(i for i in x)[:1:1] 15 15 | list(i for i in x)[:1:2] -RUF015.py:13:1: RUF015 [*] Prefer `[next(iter(x))]` over `list(x)[:1]` +RUF015.py:13:1: RUF015 [*] Prefer `[next(iter(x))]` over single element slice | 11 | tuple(x)[:1:2] 12 | list(i for i in x)[0] @@ -210,7 +210,7 @@ RUF015.py:13:1: RUF015 [*] Prefer `[next(iter(x))]` over `list(x)[:1]` 15 15 | list(i for i in x)[:1:2] 16 16 | [i for i in x][0] -RUF015.py:14:1: RUF015 [*] Prefer `[next(iter(x))]` over `list(x)[:1]` +RUF015.py:14:1: RUF015 [*] Prefer `[next(iter(x))]` over single element slice | 12 | list(i for i in x)[0] 13 | list(i for i in x)[:1] @@ -231,7 +231,7 @@ RUF015.py:14:1: RUF015 [*] Prefer `[next(iter(x))]` over `list(x)[:1]` 16 16 | [i for i in x][0] 17 17 | [i for i in x][:1] -RUF015.py:15:1: RUF015 [*] Prefer `[next(iter(x))]` over `list(x)[:1]` +RUF015.py:15:1: RUF015 [*] Prefer `[next(iter(x))]` over single element slice | 13 | list(i for i in x)[:1] 14 | list(i for i in x)[:1:1] @@ -252,7 +252,7 @@ RUF015.py:15:1: RUF015 [*] Prefer `[next(iter(x))]` over `list(x)[:1]` 17 17 | [i for i in x][:1] 18 18 | [i for i in x][:1:1] -RUF015.py:16:1: RUF015 [*] Prefer `next(iter(x))` over `list(x)[0]` +RUF015.py:16:1: RUF015 [*] Prefer `next(iter(x))` over single element slice | 14 | list(i for i in x)[:1:1] 15 | list(i for i in x)[:1:2] @@ -273,7 +273,7 @@ RUF015.py:16:1: RUF015 [*] Prefer `next(iter(x))` over `list(x)[0]` 18 18 | [i for i in x][:1:1] 19 19 | [i for i in x][:1:2] -RUF015.py:17:1: RUF015 [*] Prefer `[next(iter(x))]` over `list(x)[:1]` +RUF015.py:17:1: RUF015 [*] Prefer `[next(iter(x))]` over single element slice | 15 | list(i for i in x)[:1:2] 16 | [i for i in x][0] @@ -294,7 +294,7 @@ RUF015.py:17:1: RUF015 [*] Prefer `[next(iter(x))]` over `list(x)[:1]` 19 19 | [i for i in x][:1:2] 20 20 | -RUF015.py:18:1: RUF015 [*] Prefer `[next(iter(x))]` over `list(x)[:1]` +RUF015.py:18:1: RUF015 [*] Prefer `[next(iter(x))]` over single element slice | 16 | [i for i in x][0] 17 | [i for i in x][:1] @@ -314,7 +314,7 @@ RUF015.py:18:1: RUF015 [*] Prefer `[next(iter(x))]` over `list(x)[:1]` 20 20 | 21 21 | # OK (not indexing (solely) the first element) -RUF015.py:19:1: RUF015 [*] Prefer `[next(iter(x))]` over `list(x)[:1]` +RUF015.py:19:1: RUF015 [*] Prefer `[next(iter(x))]` over single element slice | 17 | [i for i in x][:1] 18 | [i for i in x][:1:1] @@ -335,7 +335,7 @@ 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]` +RUF015.py:38:1: RUF015 [*] Prefer `next(i + 1 for i in x)` over single element slice | 37 | # RUF015 (doesn't mirror the underlying list) 38 | [i + 1 for i in x][0] @@ -343,19 +343,19 @@ RUF015.py:38:1: RUF015 [*] Prefer `next(iter(i + 1 for i in x))` over `list(i + 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))` + = help: Replace with `next(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)) + 38 |+next(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]` +RUF015.py:39:1: RUF015 [*] Prefer `next(i for i in x if i > 5)` over single element slice | 37 | # RUF015 (doesn't mirror the underlying list) 38 | [i + 1 for i in x][0] @@ -363,19 +363,19 @@ RUF015.py:39:1: RUF015 [*] Prefer `next(iter(i for i in x if i > 5))` over `list | ^^^^^^^^^^^^^^^^^^^^^^^^^^ RUF015 40 | [(i, i + 1) for i in x][0] | - = help: Replace with `next(iter(i for i in x if i > 5))` + = help: Replace with `next(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)) + 39 |+next(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]` +RUF015.py:40:1: RUF015 [*] Prefer `next((i, i + 1) for i in x)` over single element slice | 38 | [i + 1 for i in x][0] 39 | [i for i in x if i > 5][0] @@ -384,19 +384,19 @@ RUF015.py:40:1: RUF015 [*] Prefer `next(iter((i, i + 1) for i in x))` over `list 41 | 42 | # RUF015 (multiple generators) | - = help: Replace with `next(iter((i, i + 1) for i in x))` + = help: Replace with `next((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)) + 40 |+next((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]` +RUF015.py:44:1: RUF015 [*] Prefer `next(i + j for i in x for j in y)` over single element slice | 42 | # RUF015 (multiple generators) 43 | y = range(10) @@ -405,19 +405,19 @@ RUF015.py:44:1: RUF015 [*] Prefer `next(iter(i + j for i in x for j in y))` over 45 | 46 | # RUF015 | - = help: Replace with `next(iter(i + j for i in x for j in y))` + = help: Replace with `next(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)) + 44 |+next(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]` +RUF015.py:47:1: RUF015 [*] Prefer `next(iter(range(10)))` over single element slice | 46 | # RUF015 47 | list(range(10))[0] @@ -437,7 +437,7 @@ RUF015.py:47:1: RUF015 [*] Prefer `next(iter(range(10)))` over `list(range(10))[ 49 49 | list(x["y"])[0] 50 50 | -RUF015.py:48:1: RUF015 [*] Prefer `next(iter(x.y))` over `list(x.y)[0]` +RUF015.py:48:1: RUF015 [*] Prefer `next(iter(x.y))` over single element slice | 46 | # RUF015 47 | list(range(10))[0] @@ -457,7 +457,7 @@ RUF015.py:48:1: RUF015 [*] Prefer `next(iter(x.y))` over `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]` +RUF015.py:49:1: RUF015 [*] Prefer `next(iter(x["y"]))` over single element slice | 47 | list(range(10))[0] 48 | list(x.y)[0] @@ -478,7 +478,7 @@ RUF015.py:49:1: RUF015 [*] Prefer `next(iter(x["y"]))` over `list(x["y"])[0]` 51 51 | # RUF015 (multi-line) 52 52 | revision_heads_map_ast = [ -RUF015.py:52:26: RUF015 [*] Prefer `next(iter(...))` over `list(...)[0]` +RUF015.py:52:26: RUF015 [*] Prefer `next(...)` over single element slice | 51 | # RUF015 (multi-line) 52 | revision_heads_map_ast = [ @@ -489,18 +489,18 @@ RUF015.py:52:26: RUF015 [*] Prefer `next(iter(...))` over `list(...)[0]` 56 | | ][0] | |____^ RUF015 | - = help: Replace with `next(iter(...))` + = help: Replace with `next(...)` ℹ 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( + 52 |+revision_heads_map_ast = next( 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 |+)) + 56 |+)