[ruff] itertools.starmap(..., zip(...)) (RUF058) (#15483)

Co-authored-by: Micha Reiser <[email protected]>

crates/ruff_linter/resources/test/fixtures/ruff/RUF058.py

@@ -0,0 +1,73 @@
+from itertools import starmap
+import itertools
+
+
+# Errors
+
+starmap(func, zip())
+starmap(func, zip([]))
+starmap(func, zip(*args))
+
+starmap(func, zip(a, b, c,),)
+
+
+starmap(
+    func,  # Foo
+    zip(
+        # Foo
+
+    )
+)
+
+(  # Foo
+    itertools
+    ) . starmap (
+
+    func,  zip (
+        a, b, c,
+    ),
+)
+
+(  # Foobar
+  (  starmap  )
+    # Bazqux
+) \
+(func,
+ ( (
+     (  # Zip
+    (
+        (  zip
+           # Zip
+           )
+    )
+  )
+     (a,  # A
+         b,  # B
+         c,  # C
+   )      )
+   ),
+)
+
+starmap(
+    func \
+    , \
+    zip \
+        (
+            a,\
+            b,\
+            c,\
+        )
+)
+
+
+# No errors
+
+starmap(func)
+starmap(func, zip(a, b, c, **kwargs))
+starmap(func, zip(a, b, c), foo)
+starmap(func, zip(a, b, c, lorem=ipsum))
+starmap(func, zip(a, b, c), lorem=ipsum)
+
+starmap(func, zip(a, b, c, strict=True))
+starmap(func, zip(a, b, c, strict=False))
+starmap(func, zip(a, b, c, strict=strict))

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

@@ -1126,6 +1126,9 @@ pub(crate) fn expression(expr: &Expr, checker: &mut Checker) {
             if checker.enabled(Rule::UnnecessaryEmptyIterableWithinDequeCall) {
                 ruff::rules::unnecessary_literal_within_deque_call(checker, call);
             }
+            if checker.enabled(Rule::StarmapZip) {
+                ruff::rules::starmap_zip(checker, call);
+            }
         }
         Expr::Dict(dict) => {
             if checker.any_enabled(&[

crates/ruff_linter/src/codes.rs

@@ -1003,6 +1003,7 @@ pub fn code_to_rule(linter: Linter, code: &str) -> Option<(RuleGroup, Rule)> {
         (Ruff, "055") => (RuleGroup::Preview, rules::ruff::rules::UnnecessaryRegularExpression),
         (Ruff, "056") => (RuleGroup::Preview, rules::ruff::rules::FalsyDictGetFallback),
         (Ruff, "057") => (RuleGroup::Preview, rules::ruff::rules::UnnecessaryRound),
+        (Ruff, "058") => (RuleGroup::Preview, rules::ruff::rules::StarmapZip),
         (Ruff, "100") => (RuleGroup::Stable, rules::ruff::rules::UnusedNOQA),
         (Ruff, "101") => (RuleGroup::Stable, rules::ruff::rules::RedirectedNOQA),
 

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

@@ -29,34 +29,15 @@ use crate::importer::ImportRequest;
 ///
 /// ## 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)
-/// )
+/// all(predicate(a, b) for a, b in some_iterable)
 /// ```
 ///
 /// 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)))
+/// all(starmap(predicate, some_iterable))
 /// ```
 ///
 /// ## References
@@ -83,7 +64,7 @@ impl Violation for ReimplementedStarmap {
 /// FURB140
 pub(crate) fn reimplemented_starmap(checker: &mut Checker, target: &StarmapCandidate) {
     // Generator should have exactly one comprehension.
-    let [comprehension @ ast::Comprehension { .. }] = target.generators() else {
+    let [comprehension] = target.generators() else {
         return;
     };
 

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

@@ -422,6 +422,7 @@ mod tests {
     #[test_case(Rule::PytestRaisesAmbiguousPattern, Path::new("RUF043.py"))]
     #[test_case(Rule::UnnecessaryRound, Path::new("RUF057.py"))]
     #[test_case(Rule::DataclassEnum, Path::new("RUF049.py"))]
+    #[test_case(Rule::StarmapZip, Path::new("RUF058.py"))]
     fn preview_rules(rule_code: Rule, path: &Path) -> Result<()> {
         let snapshot = format!(
             "preview__{}_{}",

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

@@ -31,6 +31,7 @@ pub(crate) use redirected_noqa::*;
 pub(crate) use redundant_bool_literal::*;
 pub(crate) use sort_dunder_all::*;
 pub(crate) use sort_dunder_slots::*;
+pub(crate) use starmap_zip::*;
 pub(crate) use static_key_dict_comprehension::*;
 #[cfg(any(feature = "test-rules", test))]
 pub(crate) use test_rules::*;
@@ -85,6 +86,7 @@ mod redundant_bool_literal;
 mod sequence_sorting;
 mod sort_dunder_all;
 mod sort_dunder_slots;
+mod starmap_zip;
 mod static_key_dict_comprehension;
 mod suppression_comment_visitor;
 #[cfg(any(feature = "test-rules", test))]

crates/ruff_linter/src/rules/ruff/rules/starmap_zip.rs

@@ -0,0 +1,146 @@
+use crate::checkers::ast::Checker;
+use ruff_diagnostics::{AlwaysFixableViolation, Applicability, Diagnostic, Edit, Fix};
+use ruff_macros::{derive_message_formats, ViolationMetadata};
+use ruff_python_ast::{parenthesize::parenthesized_range, Expr, ExprCall};
+use ruff_python_parser::TokenKind;
+use ruff_text_size::{Ranged, TextRange};
+
+/// ## What it does
+/// Checks for `itertools.starmap` calls where the second argument is a `zip` call.
+///
+/// ## Why is this bad?
+/// `zip`-ping iterables only to unpack them later from within `starmap` is unnecessary.
+/// For such cases, `map()` should be used instead.
+///
+/// ## Example
+///
+/// ```python
+/// from itertools import starmap
+///
+///
+/// starmap(func, zip(a, b))
+/// starmap(func, zip(a, b, strict=True))
+/// ```
+///
+/// Use instead:
+///
+/// ```python
+/// map(func, a, b)
+/// map(func, a, b, strict=True)  # 3.14+
+/// ```
+#[derive(ViolationMetadata)]
+pub(crate) struct StarmapZip;
+
+impl AlwaysFixableViolation for StarmapZip {
+    #[derive_message_formats]
+    fn message(&self) -> String {
+        "`itertools.starmap` called on `zip` iterable".to_string()
+    }
+
+    fn fix_title(&self) -> String {
+        "Use `map` instead".to_string()
+    }
+}
+
+/// RUF058
+pub(crate) fn starmap_zip(checker: &mut Checker, call: &ExprCall) {
+    let semantic = checker.semantic();
+
+    if !call.arguments.keywords.is_empty() {
+        return;
+    }
+
+    let [_map_func, Expr::Call(iterable_call)] = &*call.arguments.args else {
+        return;
+    };
+
+    if !iterable_call.arguments.keywords.is_empty() {
+        // TODO: Pass `strict=` to `map` too when 3.14 is supported.
+        return;
+    }
+
+    if !semantic
+        .resolve_qualified_name(&call.func)
+        .is_some_and(|it| matches!(it.segments(), ["itertools", "starmap"]))
+    {
+        return;
+    }
+
+    if !semantic.match_builtin_expr(&iterable_call.func, "zip") {
+        return;
+    }
+
+    let fix = replace_with_map(call, iterable_call, checker);
+    let diagnostic = Diagnostic::new(StarmapZip, call.range);
+
+    checker.diagnostics.push(diagnostic.with_fix(fix));
+}
+
+fn replace_with_map(starmap: &ExprCall, zip: &ExprCall, checker: &Checker) -> Fix {
+    let change_func_to_map = Edit::range_replacement("map".to_string(), starmap.func.range());
+
+    let mut remove_zip = vec![];
+
+    let full_zip_range = parenthesized_range(
+        zip.into(),
+        starmap.into(),
+        checker.comment_ranges(),
+        checker.source(),
+    )
+    .unwrap_or(zip.range());
+
+    // Delete any parentheses around the `zip` call to prevent that the argument turns into a tuple.
+    remove_zip.push(Edit::range_deletion(TextRange::new(
+        full_zip_range.start(),
+        zip.start(),
+    )));
+
+    let full_zip_func_range = parenthesized_range(
+        (&zip.func).into(),
+        zip.into(),
+        checker.comment_ranges(),
+        checker.source(),
+    )
+    .unwrap_or(zip.func.range());
+
+    // Delete the `zip` callee
+    remove_zip.push(Edit::range_deletion(full_zip_func_range));
+
+    // Delete the `(` from the `zip(...)` call
+    remove_zip.push(Edit::range_deletion(zip.arguments.l_paren_range()));
+
+    // `zip` can be called without arguments but `map` can't.
+    if zip.arguments.is_empty() {
+        remove_zip.push(Edit::insertion("[]".to_string(), zip.arguments.start()));
+    }
+
+    let after_zip = checker.tokens().after(full_zip_range.end());
+
+    // Remove any trailing commas after the `zip` call to avoid multiple trailing commas
+    // if the iterable has a trailing comma.
+    if let Some(trailing_comma) = after_zip.iter().find(|token| !token.kind().is_trivia()) {
+        if trailing_comma.kind() == TokenKind::Comma {
+            remove_zip.push(Edit::range_deletion(trailing_comma.range()));
+        }
+    }
+
+    // Delete the `)` from the `zip(...)` call
+    remove_zip.push(Edit::range_deletion(zip.arguments.r_paren_range()));
+
+    // Delete any trailing parentheses wrapping the `zip` call.
+    remove_zip.push(Edit::range_deletion(TextRange::new(
+        zip.end(),
+        full_zip_range.end(),
+    )));
+
+    let comment_ranges = checker.comment_ranges();
+    let applicability = if comment_ranges.intersects(starmap.func.range())
+        || comment_ranges.intersects(full_zip_range)
+    {
+        Applicability::Unsafe
+    } else {
+        Applicability::Safe
+    };
+
+    Fix::applicable_edits(change_func_to_map, remove_zip, applicability)
+}