diff --git a/crates/ruff/src/rules/ruff/rules/unnecessary_list_allocation_for_first_element.rs b/crates/ruff/src/rules/ruff/rules/unnecessary_list_allocation_for_first_element.rs index 886da94a551d9b..c21b43de7e853d 100644 --- a/crates/ruff/src/rules/ruff/rules/unnecessary_list_allocation_for_first_element.rs +++ b/crates/ruff/src/rules/ruff/rules/unnecessary_list_allocation_for_first_element.rs @@ -1,6 +1,7 @@ use num_traits::ToPrimitive; use ruff_diagnostics::{AlwaysAutofixableViolation, Diagnostic, Edit, Fix}; use ruff_macros::{derive_message_formats, violation}; +use ruff_python_ast::helpers::contains_effect; use rustpython_parser::ast::{self, Constant, Expr}; use crate::checkers::ast::Checker; @@ -67,7 +68,9 @@ pub(crate) fn unnecessary_list_allocation_for_first_element( let Expr::Subscript(ast::ExprSubscript { value, slice, range, .. }) = subscript else { return; }; - if !indexes_first_element(slice) { + + let (indexes_first_element, in_slice) = classify_subscript(slice); + if !indexes_first_element { return; } let Some(iter_name) = get_iterable_name(checker, value) else { @@ -80,7 +83,12 @@ pub(crate) fn unnecessary_list_allocation_for_first_element( ); if checker.patch(diagnostic.kind.rule()) { - let replacement = format!("next(iter({iter_name}))"); + let replacement = if in_slice { + format!("[next(iter({iter_name}))]") + } else { + format!("next(iter({iter_name}))") + }; + diagnostic.set_fix(Fix::suggested(Edit::range_replacement(replacement, *range))); } @@ -130,11 +138,14 @@ fn get_iterable_name<'a>(checker: &mut Checker, expr: &'a Expr) -> Option<&'a st } } -/// Check that the slice [`Expr`] is functionally equivalent to slicing into the first element. -fn indexes_first_element(expr: &Expr) -> bool { +/// Check that the slice [`Expr`] is functionally equivalent to slicing into the first element. The +/// first `bool` checks that the element is in fact first, the second checks if it's a slice or an +/// index. +fn classify_subscript(expr: &Expr) -> (bool, bool) { match expr { Expr::Constant(ast::ExprConstant { .. }) => { - get_effective_index(expr).unwrap_or(0i64) == 0i64 + let effective_index = get_effective_index(expr); + (acts_as_zero(effective_index), false) } Expr::Slice(ast::ExprSlice { step: step_value, @@ -142,26 +153,22 @@ fn indexes_first_element(expr: &Expr) -> bool { 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; + let lower = lower_index.as_ref().and_then(|l| get_effective_index(l)); + let upper = upper_index.as_ref().and_then(|u| get_effective_index(u)); + let step = step_value.as_ref().and_then(|s| get_effective_index(s)); + + let in_slice = upper.is_some() || step.is_some(); + if acts_as_zero(lower) { + if upper.unwrap_or(i64::MAX) > step.unwrap_or(1i64) { + return (false, in_slice); + } + + return (true, in_slice); } - false + (false, in_slice) } - _ => false, + _ => (false, false), } } @@ -181,3 +188,7 @@ fn get_name_id(expr: &Expr) -> Option<&str> { _ => None, } } + +fn acts_as_zero(i: Option) -> bool { + i.is_none() || i == Some(0i64) +} 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 03b99faa4f7175..0f0cdf05b2f7f7 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 @@ -37,7 +37,7 @@ RUF015.py:5:1: RUF015 [*] Prefer `next(iter(x))` over `list(x)[0]` or equivalent 3 3 | # RUF015 4 4 | list(x)[0] 5 |-list(x)[:1] - 5 |+next(iter(x)) + 5 |+[next(iter(x))] 6 6 | list(x)[:1:1] 7 7 | list(x)[:1:2] 8 8 | [i for i in x][0] @@ -58,7 +58,7 @@ RUF015.py:6:1: RUF015 [*] Prefer `next(iter(x))` over `list(x)[0]` or equivalent 4 4 | list(x)[0] 5 5 | list(x)[:1] 6 |-list(x)[:1:1] - 6 |+next(iter(x)) + 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] @@ -79,7 +79,7 @@ RUF015.py:7:1: RUF015 [*] Prefer `next(iter(x))` over `list(x)[0]` or equivalent 5 5 | list(x)[:1] 6 6 | list(x)[:1:1] 7 |-list(x)[:1:2] - 7 |+next(iter(x)) + 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] @@ -121,7 +121,7 @@ RUF015.py:9:1: RUF015 [*] Prefer `next(iter(x))` over `list(x)[0]` or equivalent 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)) + 9 |+[next(iter(x))] 10 10 | [i for i in x][:1:1] 11 11 | [i for i in x][:1:2] 12 12 | @@ -141,7 +141,7 @@ RUF015.py:10:1: RUF015 [*] Prefer `next(iter(x))` over `list(x)[0]` or equivalen 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)) + 10 |+[next(iter(x))] 11 11 | [i for i in x][:1:2] 12 12 | 13 13 | # Fine - not indexing (solely) the first element @@ -162,7 +162,7 @@ RUF015.py:11:1: RUF015 [*] Prefer `next(iter(x))` over `list(x)[0]` or equivalen 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)) + 11 |+[next(iter(x))] 12 12 | 13 13 | # Fine - not indexing (solely) the first element 14 14 | list(x)