Genericize generator name extraction

crates/ruff/resources/test/fixtures/ruff/RUF015.py

@@ -5,6 +5,14 @@
 list(x)[:1]
 list(x)[:1:1]
 list(x)[:1:2]
+tuple(x)[0]
+tuple(x)[:1]
+tuple(x)[:1:1]
+tuple(x)[:1:2]
+list(i for i in x)[0]
+list(i for i in x)[:1]
+list(i for i in x)[:1:1]
+list(i for i in x)[:1:2]
 [i for i in x][0]
 [i for i in x][:1]
 [i for i in x][:1:1]
@@ -34,4 +42,3 @@
 y = range(10)
 # Fine - multiple generators
 [i + j for i in x for j in y][0]
-

crates/ruff/src/checkers/ast/mod.rs

@@ -2203,8 +2203,8 @@ where
                 if self.enabled(Rule::UncapitalizedEnvironmentVariables) {
                     flake8_simplify::rules::use_capital_environment_variables(self, expr);
                 }
-                if self.enabled(Rule::UnnecessaryListAllocationForFirstElement) {
-                    ruff::rules::unnecessary_list_allocation_for_first_element(self, subscript);
+                if self.enabled(Rule::UnnecessaryIterableAllocationForFirstElement) {
+                    ruff::rules::unnecessary_iterable_allocation_for_first_element(self, subscript);
                 }
 
                 pandas_vet::rules::subscript(self, value, expr);

crates/ruff/src/codes.rs

@@ -777,7 +777,7 @@ pub fn code_to_rule(linter: Linter, code: &str) -> Option<(RuleGroup, Rule)> {
         (Ruff, "013") => (RuleGroup::Unspecified, rules::ruff::rules::ImplicitOptional),
         #[cfg(feature = "unreachable-code")]
         (Ruff, "014") => (RuleGroup::Nursery, rules::ruff::rules::UnreachableCode),
-        (Ruff, "015") => (RuleGroup::Unspecified, rules::ruff::rules::UnnecessaryListAllocationForFirstElement),
+        (Ruff, "015") => (RuleGroup::Unspecified, rules::ruff::rules::UnnecessaryIterableAllocationForFirstElement),
         (Ruff, "100") => (RuleGroup::Unspecified, rules::ruff::rules::UnusedNOQA),
         (Ruff, "200") => (RuleGroup::Unspecified, rules::ruff::rules::InvalidPyprojectToml),
 

crates/ruff/src/rules/ruff/mod.rs

@@ -30,7 +30,10 @@ mod tests {
     #[test_case(Rule::MutableDataclassDefault, Path::new("RUF008.py"))]
     #[test_case(Rule::PairwiseOverZipped, Path::new("RUF007.py"))]
     #[test_case(Rule::StaticKeyDictComprehension, Path::new("RUF011.py"))]
-    #[test_case(Rule::UnnecessaryListAllocationForFirstElement, Path::new("RUF015.py"))]
+    #[test_case(
+        Rule::UnnecessaryIterableAllocationForFirstElement,
+        Path::new("RUF015.py")
+    )]
     #[cfg_attr(
         feature = "unreachable-code",
         test_case(Rule::UnreachableCode, Path::new("RUF014.py"))

crates/ruff/src/rules/ruff/rules/mod.rs

@@ -9,7 +9,7 @@ pub(crate) use mutable_class_default::*;
 pub(crate) use mutable_dataclass_default::*;
 pub(crate) use pairwise_over_zipped::*;
 pub(crate) use static_key_dict_comprehension::*;
-pub(crate) use unnecessary_list_allocation_for_first_element::*;
+pub(crate) use unnecessary_iterable_allocation_for_first_element::*;
 #[cfg(feature = "unreachable-code")]
 pub(crate) use unreachable::*;
 pub(crate) use unused_noqa::*;
@@ -27,7 +27,7 @@ mod mutable_class_default;
 mod mutable_dataclass_default;
 mod pairwise_over_zipped;
 mod static_key_dict_comprehension;
-mod unnecessary_list_allocation_for_first_element;
+mod unnecessary_iterable_allocation_for_first_element;
 #[cfg(feature = "unreachable-code")]
 pub(crate) mod unreachable;
 mod unused_noqa;

crates/ruff/src/rules/ruff/rules/unnecessary_list_allocation_for_first_element.rs → crates/ruff/src/rules/ruff/rules/unnecessary_iterable_allocation_for_first_element.rs

@@ -1,7 +1,7 @@
 use num_traits::ToPrimitive;
 use ruff_diagnostics::{AlwaysAutofixableViolation, Diagnostic, Edit, Fix};
 use ruff_macros::{derive_message_formats, violation};
-use rustpython_parser::ast::{self, Constant, Expr};
+use rustpython_parser::ast::{self, Comprehension, Constant, Expr};
 
 use crate::checkers::ast::Checker;
 use crate::registry::AsRule;
@@ -32,17 +32,17 @@ use crate::registry::AsRule;
 /// - [Iterators and Iterables in Python: Run Efficient
 /// Iterations](https://realpython.com/python-iterators-iterables/#when-to-use-an-iterator-in-python)
 #[violation]
-pub(crate) struct UnnecessaryListAllocationForFirstElement {
+pub(crate) struct UnnecessaryIterableAllocationForFirstElement {
     arg: String,
 }
 
-impl UnnecessaryListAllocationForFirstElement {
+impl UnnecessaryIterableAllocationForFirstElement {
     pub(crate) fn new(arg: String) -> Self {
         Self { arg }
     }
 }
 
-impl AlwaysAutofixableViolation for UnnecessaryListAllocationForFirstElement {
+impl AlwaysAutofixableViolation for UnnecessaryIterableAllocationForFirstElement {
     #[derive_message_formats]
     fn message(&self) -> String {
         format!(
@@ -52,37 +52,56 @@ impl AlwaysAutofixableViolation for UnnecessaryListAllocationForFirstElement {
     }
 
     fn autofix_title(&self) -> String {
-        format!(
-            "Convert `list({})[0]` or equivalent list comprehension call to `next(iter({}))`",
-            self.arg, self.arg
-        )
+        format!("Replace with `next(iter({}))`", self.arg)
+    }
+}
+
+/// Contains information about a [`Expr::Subscript`] to determine if and how it accesses the first element
+/// of an iterable.
+struct ClassifiedSubscript {
+    /// If the subscript accesses the first element of the iterable
+    indexes_first_element: bool,
+    /// If the subscript is a slice (e.g., `[:1]`)
+    is_slice: bool,
+}
+
+impl ClassifiedSubscript {
+    fn new(indexes_first_element: bool, is_slice: bool) -> Self {
+        Self {
+            indexes_first_element,
+            is_slice,
+        }
     }
 }
 
 /// RUF015
-pub(crate) fn unnecessary_list_allocation_for_first_element(
+pub(crate) fn unnecessary_iterable_allocation_for_first_element(
     checker: &mut Checker,
     subscript: &Expr,
 ) {
     let Expr::Subscript(ast::ExprSubscript { value, slice, range, .. }) = subscript else {
         return;
     };
 
-    let (indexes_first_element, in_slice) = classify_subscript(slice);
+    let ClassifiedSubscript {
+        indexes_first_element,
+        is_slice,
+    } = classify_subscript(slice);
     if !indexes_first_element {
         return;
     }
+
     let Some(iter_name) = get_iterable_name(checker, value) else {
         return;
     };
 
     let mut diagnostic = Diagnostic::new(
-        UnnecessaryListAllocationForFirstElement::new(iter_name.to_string()),
+        UnnecessaryIterableAllocationForFirstElement::new(iter_name.to_string()),
         *range,
     );
 
     if checker.patch(diagnostic.kind.rule()) {
-        let replacement = if in_slice {
+        let replacement = if is_slice {
             format!("[next(iter({iter_name}))]")
         } else {
             format!("next(iter({iter_name}))")
@@ -94,62 +113,14 @@ pub(crate) fn unnecessary_list_allocation_for_first_element(
     checker.diagnostics.push(diagnostic);
 }
 
-/// Fetch the name of the iterable from a list expression if the expression returns an unmodified list
-/// which can be sliced into.
-fn get_iterable_name<'a>(checker: &mut Checker, expr: &'a Expr) -> Option<&'a str> {
-    match expr {
-        Expr::Call(ast::ExprCall { func, args, .. }) => {
-            let Some(id) = get_name_id(func.as_ref()) else {
-                return None;
-            };
-            if !(id == "list" && checker.semantic().is_builtin("list")) {
-                return None;
-            }
-
-            let Some(Expr::Name(ast::ExprName { id: arg_name, .. })) = args.first() else {
-                return None;
-            };
-
-            Some(arg_name.as_str())
-        }
-        Expr::ListComp(ast::ExprListComp {
-            elt, generators, ..
-        }) => {
-            // 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 - for example, `[i + 1 for i in x][0]`.
-            if !matches!(elt.as_ref(), Expr::Name(ast::ExprName { .. })) {
-                return None;
-            }
-
-            // If there's more than 1 generator, we can't safely say that it fits the diagnostic conditions -
-            // for example, `[(i, j) for i in x for j in y][0]`.
-            if generators.len() != 1 {
-                return None;
-            }
-
-            let generator = &generators[0];
-            // Ignore if there's an `if` statement in the comprehension, since it filters the list.
-            if !generator.ifs.is_empty() {
-                return None;
-            }
-
-            let Some(arg_name) = get_name_id(&generator.iter) else {
-                return None;
-            };
-            Some(arg_name)
-        }
-        _ => None,
-    }
-}
-
 /// 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) {
+fn classify_subscript(expr: &Expr) -> ClassifiedSubscript {
     match expr {
         Expr::Constant(ast::ExprConstant { .. }) => {
             let effective_index = get_effective_index(expr);
-            (acts_as_zero(effective_index), false)
+            ClassifiedSubscript::new(matches!(effective_index, None | Some(0)), false)
         }
         Expr::Slice(ast::ExprSlice {
             step: step_value,
@@ -161,38 +132,92 @@ fn classify_subscript(expr: &Expr) -> (bool, bool) {
             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) {
+            let is_slice = upper.is_some() || step.is_some();
+            if matches!(lower, None | Some(0)) {
                 if upper.unwrap_or(i64::MAX) > step.unwrap_or(1i64) {
-                    return (false, in_slice);
+                    return ClassifiedSubscript::new(false, is_slice);
                 }
 
-                return (true, in_slice);
+                return ClassifiedSubscript::new(true, is_slice);
             }
 
-            (false, in_slice)
+            ClassifiedSubscript::new(false, is_slice)
         }
-        _ => (false, false),
+        _ => ClassifiedSubscript::new(false, false),
     }
 }
 
-fn get_effective_index(expr: &Expr) -> Option<i64> {
+/// Fetch the name of the iterable from an expression if the expression returns an unmodified list
+/// which can be sliced into.
+fn get_iterable_name<'a>(checker: &mut Checker, expr: &'a Expr) -> Option<&'a str> {
     match expr {
-        Expr::Constant(ast::ExprConstant {
-            value: Constant::Int(value),
-            ..
-        }) => value.to_i64(),
+        Expr::Call(ast::ExprCall { func, args, .. }) => {
+            let Some(id) = get_name_id(func.as_ref()) else {
+                return None;
+            };
+            if !((id == "list" && checker.semantic().is_builtin("list"))
+                || (id == "tuple" && checker.semantic().is_builtin("tuple")))
+            {
+                return None;
+            }
+
+            match args.first() {
+                Some(Expr::Name(ast::ExprName { id: arg_name, .. })) => Some(arg_name.as_str()),
+                Some(Expr::GeneratorExp(ast::ExprGeneratorExp {
+                    elt, generators, ..
+                })) => get_generator_iterable(elt, generators),
+                _ => None,
+            }
+        }
+        Expr::ListComp(ast::ExprListComp {
+            elt, generators, ..
+        }) => get_generator_iterable(elt, generators),
         _ => None,
     }
 }
 
+fn get_generator_iterable<'a>(
+    elt: &'a Expr,
+    generators: &'a Vec<Comprehension>,
+) -> 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 - for example, `[i + 1 for i in x][0]`.
+    if !matches!(elt, Expr::Name(ast::ExprName { .. })) {
+        return None;
+    }
+
+    // If there's more than 1 generator, we can't safely say that it fits the diagnostic conditions -
+    // for example, `[(i, j) for i in x for j in y][0]`.
+    if generators.len() != 1 {
+        return None;
+    }
+
+    let generator = &generators[0];
+    // Ignore if there's an `if` statement in the comprehension, since it filters the list.
+    if !generator.ifs.is_empty() {
+        return None;
+    }
+
+    let Some(arg_name) = get_name_id(&generator.iter) else {
+                return None;
+            };
+
+    Some(arg_name)
+}
+
 fn get_name_id(expr: &Expr) -> Option<&str> {
     match expr {
         Expr::Name(ast::ExprName { id, .. }) => Some(id),
         _ => None,
     }
 }
 
-fn acts_as_zero(i: Option<i64>) -> bool {
-    i.is_none() || i == Some(0i64)
+fn get_effective_index(expr: &Expr) -> Option<i64> {
+    match expr {
+        Expr::Constant(ast::ExprConstant {
+            value: Constant::Int(value),
+            ..
+        }) => value.to_i64(),
+        _ => None,
+    }
 }