Wrap fix in [] if slicing rather than indexing

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,38 +138,37 @@ 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,
             lower: lower_index,
             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<i64>) -> bool {
+    i.is_none() || i == Some(0i64)
+}

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)