Extend reimplemented-starmap (FURB140) to catch calls with a single and starred argument (#7768)

crates/ruff_linter/resources/test/fixtures/refurb/FURB140.py

@@ -24,6 +24,12 @@ from itertools import starmap as sm
 # FURB140
 {print(x, y) for x, y in zipped()}
 
+# FURB140 (check it still flags starred arguments).
+# See https://github.com/astral-sh/ruff/issues/7636
+[foo(*t) for t in [(85, 60), (100, 80)]]
+(foo(*t) for t in [(85, 60), (100, 80)])
+{foo(*t) for t in [(85, 60), (100, 80)]}
+
 # Non-errors.
 
 [print(x, int) for x, _ in zipped()]
@@ -41,3 +47,9 @@ from itertools import starmap as sm
 [print() for x, y in zipped()]
 
 [print(x, end=y) for x, y in zipped()]
+
+[print(*x, y) for x, y in zipped()]
+
+[print(x, *y) for x, y in zipped()]
+
+[print(*x, *y) for x, y in zipped()]

crates/ruff_linter/src/rules/refurb/rules/reimplemented_starmap.rs

@@ -82,7 +82,7 @@ pub(crate) fn reimplemented_starmap(checker: &mut Checker, target: &StarmapCandi
     // ```
     //
     // `x, y, z, ...` are what we call `elts` for short.
-    let Some((elts, iter)) = match_comprehension(comprehension) else {
+    let Some(value) = match_comprehension_target(comprehension) else {
         return;
     };
 
@@ -99,14 +99,31 @@ pub(crate) fn reimplemented_starmap(checker: &mut Checker, target: &StarmapCandi
         return;
     };
 
-    // Here we want to check that `args` and `elts` are the same (same length, same elements,
-    // same order).
-    if elts.len() != args.len()
-        || !std::iter::zip(elts, args)
+    match value {
+        // Ex) `f(*x) for x in iter`
+        ComprehensionTarget::Name(name) => {
+            let [arg] = args else {
+                return;
+            };
+
+            let Expr::Starred(ast::ExprStarred { value, .. }) = arg else {
+                return;
+            };
+
+            if ComparableExpr::from(value.as_ref()) != ComparableExpr::from(name) {
+                return;
+            }
+        }
+        // Ex) `f(x, y, z) for x, y, z in iter`
+        ComprehensionTarget::Tuple(tuple) => {
+            if tuple.elts.len() != args.len()
+                || !std::iter::zip(&tuple.elts, args)
                     .all(|(x, y)| ComparableExpr::from(x) == ComparableExpr::from(y))
             {
                 return;
             }
+        }
+    }
 
     let mut diagnostic = Diagnostic::new(ReimplementedStarmap, target.range());
     if checker.patch(diagnostic.kind.rule()) {
@@ -127,7 +144,7 @@ pub(crate) fn reimplemented_starmap(checker: &mut Checker, target: &StarmapCandi
             // - For list and set comprehensions, we'd want to wrap it with `list` and `set`
             //   correspondingly.
             let main_edit = Edit::range_replacement(
-                target.try_make_suggestion(starmap_name, iter, func, checker)?,
+                target.try_make_suggestion(starmap_name, &comprehension.iter, func, checker)?,
                 target.range(),
             );
             Ok(Fix::suggested_edits(import_edit, [main_edit]))
@@ -252,7 +269,7 @@ fn try_construct_call(
     // We can only do our fix if `builtin` identifier is still bound to
     // the built-in type.
     if !checker.semantic().is_builtin(builtin) {
-        bail!(format!("Can't use built-in `{builtin}` constructor"))
+        bail!("Can't use built-in `{builtin}` constructor")
     }
 
     // In general, we replace:
@@ -306,14 +323,24 @@ fn wrap_with_call_to(call: ast::ExprCall, func_name: &str) -> ast::ExprCall {
     }
 }
 
-/// Match that the given comprehension is `(x, y, z, ...) in iter`.
-fn match_comprehension(comprehension: &ast::Comprehension) -> Option<(&[Expr], &Expr)> {
+#[derive(Debug)]
+enum ComprehensionTarget<'a> {
+    /// E.g., `(x, y, z, ...)` in `(x, y, z, ...) in iter`.
+    Tuple(&'a ast::ExprTuple),
+    /// E.g., `x` in `x in iter`.
+    Name(&'a ast::ExprName),
+}
+
+/// Extract the target from the comprehension (e.g., `(x, y, z)` in `(x, y, z, ...) in iter`).
+fn match_comprehension_target(comprehension: &ast::Comprehension) -> Option<ComprehensionTarget> {
     if comprehension.is_async || !comprehension.ifs.is_empty() {
         return None;
     }
-
-    let ast::ExprTuple { elts, .. } = comprehension.target.as_tuple_expr()?;
-    Some((elts, &comprehension.iter))
+    match &comprehension.target {
+        Expr::Tuple(tuple) => Some(ComprehensionTarget::Tuple(tuple)),
+        Expr::Name(name) => Some(ComprehensionTarget::Name(name)),
+        _ => None,
+    }
 }
 
 /// Match that the given expression is `func(x, y, z, ...)`.

crates/ruff_linter/src/rules/refurb/snapshots/ruff_linter__rules__refurb__tests__FURB140_FURB140.py.snap

@@ -119,7 +119,7 @@ FURB140.py:25:1: FURB140 [*] Use `itertools.starmap` instead of the generator
 25 | {print(x, y) for x, y in zipped()}
    | ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ FURB140
 26 | 
-27 | # Non-errors.
+27 | # FURB140 (check it still flags starred arguments).
    |
    = help: Replace with `itertools.starmap`
 
@@ -130,7 +130,69 @@ FURB140.py:25:1: FURB140 [*] Use `itertools.starmap` instead of the generator
 25    |-{print(x, y) for x, y in zipped()}
    25 |+set(sm(print, zipped()))
 26 26 | 
-27 27 | # Non-errors.
-28 28 | 
+27 27 | # FURB140 (check it still flags starred arguments).
+28 28 | # See https://github.com/astral-sh/ruff/issues/7636
+
+FURB140.py:29:1: FURB140 [*] Use `itertools.starmap` instead of the generator
+   |
+27 | # FURB140 (check it still flags starred arguments).
+28 | # See https://github.com/astral-sh/ruff/issues/7636
+29 | [foo(*t) for t in [(85, 60), (100, 80)]]
+   | ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ FURB140
+30 | (foo(*t) for t in [(85, 60), (100, 80)])
+31 | {foo(*t) for t in [(85, 60), (100, 80)]}
+   |
+   = help: Replace with `itertools.starmap`
+
+ℹ Suggested fix
+26 26 | 
+27 27 | # FURB140 (check it still flags starred arguments).
+28 28 | # See https://github.com/astral-sh/ruff/issues/7636
+29    |-[foo(*t) for t in [(85, 60), (100, 80)]]
+   29 |+list(sm(foo, [(85, 60), (100, 80)]))
+30 30 | (foo(*t) for t in [(85, 60), (100, 80)])
+31 31 | {foo(*t) for t in [(85, 60), (100, 80)]}
+32 32 | 
+
+FURB140.py:30:1: FURB140 [*] Use `itertools.starmap` instead of the generator
+   |
+28 | # See https://github.com/astral-sh/ruff/issues/7636
+29 | [foo(*t) for t in [(85, 60), (100, 80)]]
+30 | (foo(*t) for t in [(85, 60), (100, 80)])
+   | ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ FURB140
+31 | {foo(*t) for t in [(85, 60), (100, 80)]}
+   |
+   = help: Replace with `itertools.starmap`
+
+ℹ Suggested fix
+27 27 | # FURB140 (check it still flags starred arguments).
+28 28 | # See https://github.com/astral-sh/ruff/issues/7636
+29 29 | [foo(*t) for t in [(85, 60), (100, 80)]]
+30    |-(foo(*t) for t in [(85, 60), (100, 80)])
+   30 |+sm(foo, [(85, 60), (100, 80)])
+31 31 | {foo(*t) for t in [(85, 60), (100, 80)]}
+32 32 | 
+33 33 | # Non-errors.
+
+FURB140.py:31:1: FURB140 [*] Use `itertools.starmap` instead of the generator
+   |
+29 | [foo(*t) for t in [(85, 60), (100, 80)]]
+30 | (foo(*t) for t in [(85, 60), (100, 80)])
+31 | {foo(*t) for t in [(85, 60), (100, 80)]}
+   | ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ FURB140
+32 | 
+33 | # Non-errors.
+   |
+   = help: Replace with `itertools.starmap`
+
+ℹ Suggested fix
+28 28 | # See https://github.com/astral-sh/ruff/issues/7636
+29 29 | [foo(*t) for t in [(85, 60), (100, 80)]]
+30 30 | (foo(*t) for t in [(85, 60), (100, 80)])
+31    |-{foo(*t) for t in [(85, 60), (100, 80)]}
+   31 |+set(sm(foo, [(85, 60), (100, 80)]))
+32 32 | 
+33 33 | # Non-errors.
+34 34 | 
 
 

crates/ruff_python_ast/src/comparable.rs

@@ -927,11 +927,7 @@ impl<'a> From<&'a ast::Expr> for ComparableExpr<'a> {
             }) => Self::Starred(ExprStarred {
                 value: value.into(),
             }),
-            ast::Expr::Name(ast::ExprName {
-                id,
-                ctx: _,
-                range: _,
-            }) => Self::Name(ExprName { id: id.as_str() }),
+            ast::Expr::Name(name) => name.into(),
             ast::Expr::List(ast::ExprList {
                 elts,
                 ctx: _,
@@ -968,6 +964,14 @@ impl<'a> From<&'a ast::Expr> for ComparableExpr<'a> {
     }
 }
 
+impl<'a> From<&'a ast::ExprName> for ComparableExpr<'a> {
+    fn from(expr: &'a ast::ExprName) -> Self {
+        Self::Name(ExprName {
+            id: expr.id.as_str(),
+        })
+    }
+}
+
 #[derive(Debug, PartialEq, Eq, Hash)]
 pub struct StmtFunctionDef<'a> {
     is_async: bool,