Add slice to message

astral-sh · Jul 10, 2023 · 95f91e1 · 95f91e1
1 parent 4497b80
commit 95f91e1
Show file tree

Hide file tree

Showing 2 changed files with 24 additions and 16 deletions.
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
@@ -34,11 +34,15 @@ use crate::registry::AsRule;
 #[violation]
 pub(crate) struct UnnecessaryIterableAllocationForFirstElement {
     arg: String,
+    contains_slice: bool,
 }
 
 impl UnnecessaryIterableAllocationForFirstElement {
-    pub(crate) fn new(arg: String) -> Self {
-        Self { arg }
+    pub(crate) fn new(arg: String, contains_slice: bool) -> Self {
+        Self {
+            arg,
+            contains_slice,
+        }
     }
 }
 
@@ -52,7 +56,11 @@ impl AlwaysAutofixableViolation for UnnecessaryIterableAllocationForFirstElement
     }
 
     fn autofix_title(&self) -> String {
-        format!("Replace with `next(iter({}))`", self.arg)
+        if self.contains_slice {
+            format!("Replace with `[next(iter({}))]", self.arg)
+        } else {
+            format!("Replace with `next(iter({}))`", self.arg)
+        }
     }
 }
 
@@ -96,7 +104,7 @@ pub(crate) fn unnecessary_iterable_allocation_for_first_element(
     };
 
     let mut diagnostic = Diagnostic::new(
-        UnnecessaryIterableAllocationForFirstElement::new(iter_name.to_string()),
+        UnnecessaryIterableAllocationForFirstElement::new(iter_name.to_string(), is_slice),
         *range,
     );
 

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
@@ -30,7 +30,7 @@ RUF015.py:5:1: RUF015 [*] Prefer `next(iter(x))` over `list(x)[0]` or equivalent
 6 | list(x)[:1:1]
 7 | list(x)[:1:2]
   |
-  = help: Replace with `next(iter(x))`
+  = help: Replace with `[next(iter(x))]
 
 ℹ Suggested fix
 2 2 | 
@@ -51,7 +51,7 @@ RUF015.py:6:1: RUF015 [*] Prefer `next(iter(x))` over `list(x)[0]` or equivalent
 7 | list(x)[:1:2]
 8 | tuple(x)[0]
   |
-  = help: Replace with `next(iter(x))`
+  = help: Replace with `[next(iter(x))]
 
 ℹ Suggested fix
 3 3 | # RUF015
@@ -72,7 +72,7 @@ RUF015.py:7:1: RUF015 [*] Prefer `next(iter(x))` over `list(x)[0]` or equivalent
 8 | tuple(x)[0]
 9 | tuple(x)[:1]
   |
-  = help: Replace with `next(iter(x))`
+  = help: Replace with `[next(iter(x))]
 
 ℹ Suggested fix
 4 4 | list(x)[0]
@@ -114,7 +114,7 @@ RUF015.py:9:1: RUF015 [*] Prefer `next(iter(x))` over `list(x)[0]` or equivalent
 10 | tuple(x)[:1:1]
 11 | tuple(x)[:1:2]
    |
-   = help: Replace with `next(iter(x))`
+   = help: Replace with `[next(iter(x))]
 
 ℹ Suggested fix
 6  6  | list(x)[:1:1]
@@ -135,7 +135,7 @@ RUF015.py:10:1: RUF015 [*] Prefer `next(iter(x))` over `list(x)[0]` or equivalen
 11 | tuple(x)[:1:2]
 12 | list(i for i in x)[0]
    |
-   = help: Replace with `next(iter(x))`
+   = help: Replace with `[next(iter(x))]
 
 ℹ Suggested fix
 7  7  | list(x)[:1:2]
@@ -156,7 +156,7 @@ RUF015.py:11:1: RUF015 [*] Prefer `next(iter(x))` over `list(x)[0]` or equivalen
 12 | list(i for i in x)[0]
 13 | list(i for i in x)[:1]
    |
-   = help: Replace with `next(iter(x))`
+   = help: Replace with `[next(iter(x))]
 
 ℹ Suggested fix
 8  8  | tuple(x)[0]
@@ -198,7 +198,7 @@ RUF015.py:13:1: RUF015 [*] Prefer `next(iter(x))` over `list(x)[0]` or equivalen
 14 | list(i for i in x)[:1:1]
 15 | list(i for i in x)[:1:2]
    |
-   = help: Replace with `next(iter(x))`
+   = help: Replace with `[next(iter(x))]
 
 ℹ Suggested fix
 10 10 | tuple(x)[:1:1]
@@ -219,7 +219,7 @@ RUF015.py:14:1: RUF015 [*] Prefer `next(iter(x))` over `list(x)[0]` or equivalen
 15 | list(i for i in x)[:1:2]
 16 | [i for i in x][0]
    |
-   = help: Replace with `next(iter(x))`
+   = help: Replace with `[next(iter(x))]
 
 ℹ Suggested fix
 11 11 | tuple(x)[:1:2]
@@ -240,7 +240,7 @@ RUF015.py:15:1: RUF015 [*] Prefer `next(iter(x))` over `list(x)[0]` or equivalen
 16 | [i for i in x][0]
 17 | [i for i in x][:1]
    |
-   = help: Replace with `next(iter(x))`
+   = help: Replace with `[next(iter(x))]
 
 ℹ Suggested fix
 12 12 | list(i for i in x)[0]
@@ -282,7 +282,7 @@ RUF015.py:17:1: RUF015 [*] Prefer `next(iter(x))` over `list(x)[0]` or equivalen
 18 | [i for i in x][:1:1]
 19 | [i for i in x][:1:2]
    |
-   = help: Replace with `next(iter(x))`
+   = help: Replace with `[next(iter(x))]
 
 ℹ Suggested fix
 14 14 | list(i for i in x)[:1:1]
@@ -302,7 +302,7 @@ RUF015.py:18:1: RUF015 [*] Prefer `next(iter(x))` over `list(x)[0]` or equivalen
    | ^^^^^^^^^^^^^^^^^^^^ RUF015
 19 | [i for i in x][:1:2]
    |
-   = help: Replace with `next(iter(x))`
+   = help: Replace with `[next(iter(x))]
 
 ℹ Suggested fix
 15 15 | list(i for i in x)[:1:2]
@@ -323,7 +323,7 @@ RUF015.py:19:1: RUF015 [*] Prefer `next(iter(x))` over `list(x)[0]` or equivalen
 20 | 
 21 | # Fine - not indexing (solely) the first element
    |
-   = help: Replace with `next(iter(x))`
+   = help: Replace with `[next(iter(x))]
 
 ℹ Suggested fix
 16 16 | [i for i in x][0]