[refurb] Implement reimplemented-starmap rule (FURB140) (#7253)

## Summary

This PR is part of a bigger effort of re-implementing `refurb` rules
#1348. It adds support for
[FURB140](https://github.com/dosisod/refurb/blob/master/refurb/checks/itertools/use_starmap.py)

## Test Plan

I included a new test + checked that all other tests pass.

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

@@ -0,0 +1,43 @@
+def zipped():
+    return zip([1, 2, 3], "ABC")
+
+# Errors.
+
+# FURB140
+[print(x, y) for x, y in zipped()]
+
+# FURB140
+(print(x, y) for x, y in zipped())
+
+# FURB140
+{print(x, y) for x, y in zipped()}
+
+
+from itertools import starmap as sm
+
+# FURB140
+[print(x, y) for x, y in zipped()]
+
+# FURB140
+(print(x, y) for x, y in zipped())
+
+# FURB140
+{print(x, y) for x, y in zipped()}
+
+# Non-errors.
+
+[print(x, int) for x, _ in zipped()]
+
+[print(x, *y) for x, y in zipped()]
+
+[print(x, y, 1) for x, y in zipped()]
+
+[print(y, x) for x, y in zipped()]
+
+[print(x + 1, y) for x, y in zipped()]
+
+[print(x) for x in range(100)]
+
+[print() for x, y in zipped()]
+
+[print(x, end=y) for x, y in zipped()]

crates/ruff/src/checkers/ast/analyze/expression.rs

@@ -1279,16 +1279,13 @@ pub(crate) fn expression(expr: &Expr, checker: &mut Checker) {
                 flake8_simplify::rules::twisted_arms_in_ifexpr(checker, expr, test, body, orelse);
             }
         }
-        Expr::ListComp(ast::ExprListComp {
+        Expr::ListComp(
+            comp @ ast::ExprListComp {
                 elt,
                 generators,
                 range: _,
-        })
-        | Expr::SetComp(ast::ExprSetComp {
-            elt,
-            generators,
-            range: _,
-        }) => {
+            },
+        ) => {
             if checker.enabled(Rule::UnnecessaryComprehension) {
                 flake8_comprehensions::rules::unnecessary_list_set_comprehension(
                     checker, expr, elt, generators,
@@ -1302,6 +1299,33 @@ pub(crate) fn expression(expr: &Expr, checker: &mut Checker) {
                     pylint::rules::iteration_over_set(checker, &generator.iter);
                 }
             }
+            if checker.enabled(Rule::ReimplementedStarmap) {
+                refurb::rules::reimplemented_starmap(checker, &comp.into());
+            }
+        }
+        Expr::SetComp(
+            comp @ ast::ExprSetComp {
+                elt,
+                generators,
+                range: _,
+            },
+        ) => {
+            if checker.enabled(Rule::UnnecessaryComprehension) {
+                flake8_comprehensions::rules::unnecessary_list_set_comprehension(
+                    checker, expr, elt, generators,
+                );
+            }
+            if checker.enabled(Rule::FunctionUsesLoopVariable) {
+                flake8_bugbear::rules::function_uses_loop_variable(checker, &Node::Expr(expr));
+            }
+            if checker.enabled(Rule::IterationOverSet) {
+                for generator in generators {
+                    pylint::rules::iteration_over_set(checker, &generator.iter);
+                }
+            }
+            if checker.enabled(Rule::ReimplementedStarmap) {
+                refurb::rules::reimplemented_starmap(checker, &comp.into());
+            }
         }
         Expr::DictComp(ast::ExprDictComp {
             key,
@@ -1326,11 +1350,13 @@ pub(crate) fn expression(expr: &Expr, checker: &mut Checker) {
                 ruff::rules::static_key_dict_comprehension(checker, key);
             }
         }
-        Expr::GeneratorExp(ast::ExprGeneratorExp {
+        Expr::GeneratorExp(
+            generator @ ast::ExprGeneratorExp {
                 generators,
                 elt: _,
                 range: _,
-        }) => {
+            },
+        ) => {
             if checker.enabled(Rule::FunctionUsesLoopVariable) {
                 flake8_bugbear::rules::function_uses_loop_variable(checker, &Node::Expr(expr));
             }
@@ -1339,6 +1365,9 @@ pub(crate) fn expression(expr: &Expr, checker: &mut Checker) {
                     pylint::rules::iteration_over_set(checker, &generator.iter);
                 }
             }
+            if checker.enabled(Rule::ReimplementedStarmap) {
+                refurb::rules::reimplemented_starmap(checker, &generator.into());
+            }
         }
         Expr::BoolOp(
             bool_op @ ast::ExprBoolOp {

crates/ruff/src/codes.rs

@@ -919,6 +919,7 @@ pub fn code_to_rule(linter: Linter, code: &str) -> Option<(RuleGroup, Rule)> {
         (Refurb, "131") => (RuleGroup::Nursery, rules::refurb::rules::DeleteFullSlice),
         #[allow(deprecated)]
         (Refurb, "132") => (RuleGroup::Nursery, rules::refurb::rules::CheckAndRemoveFromSet),
+        (Refurb, "140") => (RuleGroup::Preview, rules::refurb::rules::ReimplementedStarmap),
         (Refurb, "145") => (RuleGroup::Preview, rules::refurb::rules::SliceCopy),
 
         // flake8-logging

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

@@ -17,6 +17,7 @@ mod tests {
     #[test_case(Rule::RepeatedAppend, Path::new("FURB113.py"))]
     #[test_case(Rule::DeleteFullSlice, Path::new("FURB131.py"))]
     #[test_case(Rule::CheckAndRemoveFromSet, Path::new("FURB132.py"))]
+    #[test_case(Rule::ReimplementedStarmap, Path::new("FURB140.py"))]
     #[test_case(Rule::SliceCopy, Path::new("FURB145.py"))]
     fn rules(rule_code: Rule, path: &Path) -> Result<()> {
         let snapshot = format!("{}_{}", rule_code.noqa_code(), path.to_string_lossy());

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

@@ -1,9 +1,11 @@
 pub(crate) use check_and_remove_from_set::*;
 pub(crate) use delete_full_slice::*;
+pub(crate) use reimplemented_starmap::*;
 pub(crate) use repeated_append::*;
 pub(crate) use slice_copy::*;
 
 mod check_and_remove_from_set;
 mod delete_full_slice;
+mod reimplemented_starmap;
 mod repeated_append;
 mod slice_copy;

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

@@ -0,0 +1,332 @@
+use anyhow::{bail, Result};
+use ruff_diagnostics::{AutofixKind, Diagnostic, Edit, Fix, Violation};
+use ruff_macros::{derive_message_formats, violation};
+use ruff_python_ast::comparable::ComparableExpr;
+use ruff_python_ast::{self as ast, Expr};
+use ruff_text_size::{Ranged, TextRange};
+
+use crate::checkers::ast::Checker;
+use crate::importer::ImportRequest;
+use crate::registry::AsRule;
+
+/// ## What it does
+/// Checks for generator expressions, list and set comprehensions that can
+/// be replaced with `itertools.starmap`.
+///
+/// ## Why is this bad?
+/// When unpacking values from iterators to pass them directly to
+/// a function, prefer `itertools.starmap`.
+///
+/// Using `itertools.starmap` is more concise and readable.
+///
+/// ## Example
+/// ```python
+/// scores = [85, 100, 60]
+/// passing_scores = [60, 80, 70]
+///
+///
+/// def passed_test(score: int, passing_score: int) -> bool:
+///     return score >= passing_score
+///
+///
+/// passed_all_tests = all(
+///     passed_test(score, passing_score)
+///     for score, passing_score in zip(scores, passing_scores)
+/// )
+/// ```
+///
+/// Use instead:
+/// ```python
+/// from itertools import starmap
+///
+///
+/// scores = [85, 100, 60]
+/// passing_scores = [60, 80, 70]
+///
+///
+/// def passed_test(score: int, passing_score: int) -> bool:
+///     return score >= passing_score
+///
+///
+/// passed_all_tests = all(starmap(passed_test, zip(scores, passing_scores)))
+/// ```
+///
+/// ## References
+/// - [Python documentation: `itertools.starmap`](https://docs.python.org/3/library/itertools.html#itertools.starmap)
+#[violation]
+pub struct ReimplementedStarmap;
+
+impl Violation for ReimplementedStarmap {
+    const AUTOFIX: AutofixKind = AutofixKind::Sometimes;
+
+    #[derive_message_formats]
+    fn message(&self) -> String {
+        format!("Use `itertools.starmap` instead of the generator")
+    }
+
+    fn autofix_title(&self) -> Option<String> {
+        Some(format!("Replace with `itertools.starmap`"))
+    }
+}
+
+/// FURB140
+pub(crate) fn reimplemented_starmap(checker: &mut Checker, target: &StarmapCandidate) {
+    // Generator should have exactly one comprehension.
+    let [comprehension @ ast::Comprehension { .. }] = target.generators() else {
+        return;
+    };
+
+    // This comprehension should have a form:
+    // ```python
+    // (x, y, z, ...) in iter
+    // ```
+    //
+    // `x, y, z, ...` are what we call `elts` for short.
+    let Some((elts, iter)) = match_comprehension(comprehension) else {
+        return;
+    };
+
+    // Generator should produce one element that should look like:
+    // ```python
+    // func(a, b, c, ...)
+    // ```
+    //
+    // here we refer to `a, b, c, ...` as `args`.
+    //
+    // NOTE: `func` is not necessarily just a function name, it can be an attribute access,
+    //       or even a call itself.
+    let Some((args, func)) = match_call(target.element()) else {
+        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)
+            .all(|(x, y)| ComparableExpr::from(x) == ComparableExpr::from(y))
+    {
+        return;
+    }
+
+    let mut diagnostic = Diagnostic::new(ReimplementedStarmap, target.range());
+    if checker.patch(diagnostic.kind.rule()) {
+        diagnostic.try_set_fix(|| {
+            // Try importing `starmap` from `itertools`.
+            //
+            // It is not required to be `itertools.starmap`, though. The user might've already
+            // imported it. Maybe even under a different name. So, we should use that name
+            // for fix construction.
+            let (import_edit, starmap_name) = checker.importer().get_or_import_symbol(
+                &ImportRequest::import_from("itertools", "starmap"),
+                target.start(),
+                checker.semantic(),
+            )?;
+            // The actual fix suggestion depends on what type of expression we were looking at.
+            //
+            // - For generator expressions, we use `starmap` call directly.
+            // - 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.range(),
+            );
+            Ok(Fix::suggested_edits(import_edit, [main_edit]))
+        });
+    }
+    checker.diagnostics.push(diagnostic);
+}
+
+/// An enum for a node that can be considered a candidate for replacement with `starmap`.
+#[derive(Debug)]
+pub(crate) enum StarmapCandidate<'a> {
+    Generator(&'a ast::ExprGeneratorExp),
+    ListComp(&'a ast::ExprListComp),
+    SetComp(&'a ast::ExprSetComp),
+}
+
+impl<'a> From<&'a ast::ExprGeneratorExp> for StarmapCandidate<'a> {
+    fn from(generator: &'a ast::ExprGeneratorExp) -> Self {
+        Self::Generator(generator)
+    }
+}
+
+impl<'a> From<&'a ast::ExprListComp> for StarmapCandidate<'a> {
+    fn from(list_comp: &'a ast::ExprListComp) -> Self {
+        Self::ListComp(list_comp)
+    }
+}
+
+impl<'a> From<&'a ast::ExprSetComp> for StarmapCandidate<'a> {
+    fn from(set_comp: &'a ast::ExprSetComp) -> Self {
+        Self::SetComp(set_comp)
+    }
+}
+
+impl Ranged for StarmapCandidate<'_> {
+    fn range(&self) -> TextRange {
+        match self {
+            Self::Generator(generator) => generator.range(),
+            Self::ListComp(list_comp) => list_comp.range(),
+            Self::SetComp(set_comp) => set_comp.range(),
+        }
+    }
+}
+
+impl StarmapCandidate<'_> {
+    /// Return the generated element for the candidate.
+    pub(crate) fn element(&self) -> &Expr {
+        match self {
+            Self::Generator(generator) => generator.elt.as_ref(),
+            Self::ListComp(list_comp) => list_comp.elt.as_ref(),
+            Self::SetComp(set_comp) => set_comp.elt.as_ref(),
+        }
+    }
+
+    /// Return the generator comprehensions for the candidate.
+    pub(crate) fn generators(&self) -> &[ast::Comprehension] {
+        match self {
+            Self::Generator(generator) => generator.generators.as_slice(),
+            Self::ListComp(list_comp) => list_comp.generators.as_slice(),
+            Self::SetComp(set_comp) => set_comp.generators.as_slice(),
+        }
+    }
+
+    /// Try to produce a fix suggestion transforming this node into a call to `starmap`.
+    pub(crate) fn try_make_suggestion(
+        &self,
+        name: String,
+        iter: &Expr,
+        func: &Expr,
+        checker: &Checker,
+    ) -> Result<String> {
+        match self {
+            Self::Generator(_) => {
+                // For generator expressions, we replace:
+                // ```python
+                // (foo(...) for ... in iter)
+                // ```
+                //
+                // with:
+                // ```python
+                // itertools.starmap(foo, iter)
+                // ```
+                let call = construct_starmap_call(name, iter, func);
+                Ok(checker.generator().expr(&call.into()))
+            }
+            Self::ListComp(_) => {
+                // For list comprehensions, we replace:
+                // ```python
+                // [foo(...) for ... in iter]
+                // ```
+                //
+                // with:
+                // ```python
+                // list(itertools.starmap(foo, iter))
+                // ```
+                try_construct_call(name, iter, func, "list", checker)
+            }
+            Self::SetComp(_) => {
+                // For set comprehensions, we replace:
+                // ```python
+                // {foo(...) for ... in iter}
+                // ```
+                //
+                // with:
+                // ```python
+                // set(itertools.starmap(foo, iter))
+                // ```
+                try_construct_call(name, iter, func, "set", checker)
+            }
+        }
+    }
+}
+
+/// Try constructing the call to `itertools.starmap` and wrapping it with the given builtin.
+fn try_construct_call(
+    name: String,
+    iter: &Expr,
+    func: &Expr,
+    builtin: &str,
+    checker: &Checker,
+) -> Result<String> {
+    // 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"))
+    }
+
+    // In general, we replace:
+    // ```python
+    // foo(...) for ... in iter
+    // ```
+    //
+    // with:
+    // ```python
+    // builtin(itertools.starmap(foo, iter))
+    // ```
+    // where `builtin` is a constructor for a target collection.
+    let call = construct_starmap_call(name, iter, func);
+    let wrapped = wrap_with_call_to(call, builtin);
+    Ok(checker.generator().expr(&wrapped.into()))
+}
+
+/// Construct the call to `itertools.starmap` for suggestion.
+fn construct_starmap_call(starmap_binding: String, iter: &Expr, func: &Expr) -> ast::ExprCall {
+    let starmap = ast::ExprName {
+        id: starmap_binding,
+        ctx: ast::ExprContext::Load,
+        range: TextRange::default(),
+    };
+    ast::ExprCall {
+        func: Box::new(starmap.into()),
+        arguments: ast::Arguments {
+            args: vec![func.clone(), iter.clone()],
+            keywords: vec![],
+            range: TextRange::default(),
+        },
+        range: TextRange::default(),
+    }
+}
+
+/// Wrap given function call with yet another call.
+fn wrap_with_call_to(call: ast::ExprCall, func_name: &str) -> ast::ExprCall {
+    let name = ast::ExprName {
+        id: func_name.to_string(),
+        ctx: ast::ExprContext::Load,
+        range: TextRange::default(),
+    };
+    ast::ExprCall {
+        func: Box::new(name.into()),
+        arguments: ast::Arguments {
+            args: vec![call.into()],
+            keywords: vec![],
+            range: TextRange::default(),
+        },
+        range: TextRange::default(),
+    }
+}
+
+/// Match that the given comprehension is `(x, y, z, ...) in iter`.
+fn match_comprehension(comprehension: &ast::Comprehension) -> Option<(&[Expr], &Expr)> {
+    if comprehension.is_async || !comprehension.ifs.is_empty() {
+        return None;
+    }
+
+    let ast::ExprTuple { elts, .. } = comprehension.target.as_tuple_expr()?;
+    Some((elts, &comprehension.iter))
+}
+
+/// Match that the given expression is `func(x, y, z, ...)`.
+fn match_call(element: &Expr) -> Option<(&[Expr], &Expr)> {
+    let ast::ExprCall {
+        func,
+        arguments: ast::Arguments { args, keywords, .. },
+        ..
+    } = element.as_call_expr()?;
+
+    if !keywords.is_empty() {
+        return None;
+    }
+
+    Some((args, func))
+}

crates/ruff/src/rules/refurb/snapshots/ruff__rules__refurb__tests__FURB140_FURB140.py.snap

@@ -0,0 +1,136 @@
+---
+source: crates/ruff/src/rules/refurb/mod.rs
+---
+FURB140.py:7:1: FURB140 [*] Use `itertools.starmap` instead of the generator
+  |
+6 | # FURB140
+7 | [print(x, y) for x, y in zipped()]
+  | ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ FURB140
+8 | 
+9 | # FURB140
+  |
+  = help: Replace with `itertools.starmap`
+
+ℹ Suggested fix
+  1 |+from itertools import starmap
+1 2 | def zipped():
+2 3 |     return zip([1, 2, 3], "ABC")
+3 4 | 
+4 5 | # Errors.
+5 6 | 
+6 7 | # FURB140
+7   |-[print(x, y) for x, y in zipped()]
+  8 |+list(starmap(print, zipped()))
+8 9 | 
+9 10 | # FURB140
+10 11 | (print(x, y) for x, y in zipped())
+
+FURB140.py:10:1: FURB140 [*] Use `itertools.starmap` instead of the generator
+   |
+ 9 | # FURB140
+10 | (print(x, y) for x, y in zipped())
+   | ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ FURB140
+11 | 
+12 | # FURB140
+   |
+   = help: Replace with `itertools.starmap`
+
+ℹ Suggested fix
+   1  |+from itertools import starmap
+1  2  | def zipped():
+2  3  |     return zip([1, 2, 3], "ABC")
+3  4  | 
+--------------------------------------------------------------------------------
+7  8  | [print(x, y) for x, y in zipped()]
+8  9  | 
+9  10 | # FURB140
+10    |-(print(x, y) for x, y in zipped())
+   11 |+starmap(print, zipped())
+11 12 | 
+12 13 | # FURB140
+13 14 | {print(x, y) for x, y in zipped()}
+
+FURB140.py:13:1: FURB140 [*] Use `itertools.starmap` instead of the generator
+   |
+12 | # FURB140
+13 | {print(x, y) for x, y in zipped()}
+   | ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ FURB140
+   |
+   = help: Replace with `itertools.starmap`
+
+ℹ Suggested fix
+   1  |+from itertools import starmap
+1  2  | def zipped():
+2  3  |     return zip([1, 2, 3], "ABC")
+3  4  | 
+--------------------------------------------------------------------------------
+10 11 | (print(x, y) for x, y in zipped())
+11 12 | 
+12 13 | # FURB140
+13    |-{print(x, y) for x, y in zipped()}
+   14 |+set(starmap(print, zipped()))
+14 15 | 
+15 16 | 
+16 17 | from itertools import starmap as sm
+
+FURB140.py:19:1: FURB140 [*] Use `itertools.starmap` instead of the generator
+   |
+18 | # FURB140
+19 | [print(x, y) for x, y in zipped()]
+   | ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ FURB140
+20 | 
+21 | # FURB140
+   |
+   = help: Replace with `itertools.starmap`
+
+ℹ Suggested fix
+16 16 | from itertools import starmap as sm
+17 17 | 
+18 18 | # FURB140
+19    |-[print(x, y) for x, y in zipped()]
+   19 |+list(sm(print, zipped()))
+20 20 | 
+21 21 | # FURB140
+22 22 | (print(x, y) for x, y in zipped())
+
+FURB140.py:22:1: FURB140 [*] Use `itertools.starmap` instead of the generator
+   |
+21 | # FURB140
+22 | (print(x, y) for x, y in zipped())
+   | ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ FURB140
+23 | 
+24 | # FURB140
+   |
+   = help: Replace with `itertools.starmap`
+
+ℹ Suggested fix
+19 19 | [print(x, y) for x, y in zipped()]
+20 20 | 
+21 21 | # FURB140
+22    |-(print(x, y) for x, y in zipped())
+   22 |+sm(print, zipped())
+23 23 | 
+24 24 | # FURB140
+25 25 | {print(x, y) for x, y in zipped()}
+
+FURB140.py:25:1: FURB140 [*] Use `itertools.starmap` instead of the generator
+   |
+24 | # FURB140
+25 | {print(x, y) for x, y in zipped()}
+   | ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ FURB140
+26 | 
+27 | # Non-errors.
+   |
+   = help: Replace with `itertools.starmap`
+
+ℹ Suggested fix
+22 22 | (print(x, y) for x, y in zipped())
+23 23 | 
+24 24 | # FURB140
+25    |-{print(x, y) for x, y in zipped()}
+   25 |+set(sm(print, zipped()))
+26 26 | 
+27 27 | # Non-errors.
+28 28 | 
+
+

crates/ruff_workspace/src/configuration.rs

@@ -854,12 +854,13 @@ mod tests {
 
     const PREVIEW_RULES: &[Rule] = &[
         Rule::DirectLoggerInstantiation,
+        Rule::InvalidGetLoggerArgument,
         Rule::ManualDictComprehension,
+        Rule::ReimplementedStarmap,
         Rule::SliceCopy,
         Rule::TooManyPublicMethods,
         Rule::TooManyPublicMethods,
         Rule::UndocumentedWarn,
-        Rule::InvalidGetLoggerArgument,
     ];
 
     #[allow(clippy::needless_pass_by_value)]

ruff.schema.json

@@ -2091,6 +2091,7 @@
         "FURB131",
         "FURB132",
         "FURB14",
+        "FURB140",
         "FURB145",
         "G",
         "G0",