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

Co-authored-by: Micha Reiser <micha@reiser.io>

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,18 +29,7 @@ use crate::importer::ImportRequest;
 ///
 /// ## Example
 /// ```python
-/// scores = [85, 100, 60]
+/// all(predicate(a, b) for a, b in some_iterable)
-/// 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:
@@ -48,15 +37,7 @@ use crate::importer::ImportRequest;
 /// from itertools import starmap
 ///
 ///
-/// scores = [85, 100, 60]
+/// all(starmap(predicate, some_iterable))
-/// 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
@@ -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)
+}

crates/ruff_linter/src/rules/ruff/snapshots/ruff_linter__rules__ruff__tests__preview__RUF058_RUF058.py.snap

@@ -0,0 +1,247 @@
+---
+source: crates/ruff_linter/src/rules/ruff/mod.rs
+---
+RUF058.py:7:1: RUF058 [*] `itertools.starmap` called on `zip` iterable
+  |
+5 | # Errors
+6 |
+7 | starmap(func, zip())
+  | ^^^^^^^^^^^^^^^^^^^^ RUF058
+8 | starmap(func, zip([]))
+9 | starmap(func, zip(*args))
+  |
+  = help: Use `map` instead
+
+ℹ Safe fix
+4 4 | 
+5 5 | # Errors
+6 6 | 
+7   |-starmap(func, zip())
+  7 |+map(func, [])
+8 8 | starmap(func, zip([]))
+9 9 | starmap(func, zip(*args))
+10 10 | 
+
+RUF058.py:8:1: RUF058 [*] `itertools.starmap` called on `zip` iterable
+  |
+7 | starmap(func, zip())
+8 | starmap(func, zip([]))
+  | ^^^^^^^^^^^^^^^^^^^^^^ RUF058
+9 | starmap(func, zip(*args))
+  |
+  = help: Use `map` instead
+
+ℹ Safe fix
+5 5 | # Errors
+6 6 | 
+7 7 | starmap(func, zip())
+8   |-starmap(func, zip([]))
+  8 |+map(func, [])
+9 9 | starmap(func, zip(*args))
+10 10 | 
+11 11 | starmap(func, zip(a, b, c,),)
+
+RUF058.py:9:1: RUF058 [*] `itertools.starmap` called on `zip` iterable
+   |
+ 7 | starmap(func, zip())
+ 8 | starmap(func, zip([]))
+ 9 | starmap(func, zip(*args))
+   | ^^^^^^^^^^^^^^^^^^^^^^^^^ RUF058
+10 |
+11 | starmap(func, zip(a, b, c,),)
+   |
+   = help: Use `map` instead
+
+ℹ Safe fix
+6  6  | 
+7  7  | starmap(func, zip())
+8  8  | starmap(func, zip([]))
+9     |-starmap(func, zip(*args))
+   9  |+map(func, *args)
+10 10 | 
+11 11 | starmap(func, zip(a, b, c,),)
+12 12 | 
+
+RUF058.py:11:1: RUF058 [*] `itertools.starmap` called on `zip` iterable
+   |
+ 9 | starmap(func, zip(*args))
+10 |
+11 | starmap(func, zip(a, b, c,),)
+   | ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ RUF058
+   |
+   = help: Use `map` instead
+
+ℹ Safe fix
+8  8  | starmap(func, zip([]))
+9  9  | starmap(func, zip(*args))
+10 10 | 
+11    |-starmap(func, zip(a, b, c,),)
+   11 |+map(func, a, b, c,)
+12 12 | 
+13 13 | 
+14 14 | starmap(
+
+RUF058.py:14:1: RUF058 [*] `itertools.starmap` called on `zip` iterable
+   |
+14 | / starmap(
+15 | |     func,  # Foo
+16 | |     zip(
+17 | |         # Foo
+18 | |
+19 | |     )
+20 | | )
+   | |_^ RUF058
+21 |
+22 |   (  # Foo
+   |
+   = help: Use `map` instead
+
+ℹ Unsafe fix
+11 11 | starmap(func, zip(a, b, c,),)
+12 12 | 
+13 13 | 
+14    |-starmap(
+   14 |+map(
+15 15 |     func,  # Foo
+16    |-    zip(
+   16 |+    []
+17 17 |         # Foo
+18 18 | 
+19    |-    )
+   19 |+    
+20 20 | )
+21 21 | 
+22 22 | (  # Foo
+
+RUF058.py:22:1: RUF058 [*] `itertools.starmap` called on `zip` iterable
+   |
+20 |   )
+21 |
+22 | / (  # Foo
+23 | |     itertools
+24 | |     ) . starmap (
+25 | |
+26 | |     func,  zip (
+27 | |         a, b, c,
+28 | |     ),
+29 | | )
+   | |_^ RUF058
+30 |
+31 |   (  # Foobar
+   |
+   = help: Use `map` instead
+
+ℹ Unsafe fix
+19 19 |     )
+20 20 | )
+21 21 | 
+22    |-(  # Foo
+23    |-    itertools
+24    |-    ) . starmap (
+   22 |+map (
+25 23 | 
+26    |-    func,  zip (
+   24 |+    func,   
+27 25 |         a, b, c,
+28    |-    ),
+   26 |+    
+29 27 | )
+30 28 | 
+31 29 | (  # Foobar
+
+RUF058.py:31:1: RUF058 [*] `itertools.starmap` called on `zip` iterable
+   |
+29 |   )
+30 |
+31 | / (  # Foobar
+32 | |   (  starmap  )
+33 | |     # Bazqux
+34 | | ) \
+35 | | (func,
+36 | |  ( (
+37 | |      (  # Zip
+38 | |     (
+39 | |         (  zip
+40 | |            # Zip
+41 | |            )
+42 | |     )
+43 | |   )
+44 | |      (a,  # A
+45 | |          b,  # B
+46 | |          c,  # C
+47 | |    )      )
+48 | |    ),
+49 | | )
+   | |_^ RUF058
+50 |
+51 |   starmap(
+   |
+   = help: Use `map` instead
+
+ℹ Unsafe fix
+29 29 | )
+30 30 | 
+31 31 | (  # Foobar
+32    |-  (  starmap  )
+   32 |+  (  map  )
+33 33 |     # Bazqux
+34 34 | ) \
+35 35 | (func,
+36    |- ( (
+37    |-     (  # Zip
+38    |-    (
+39    |-        (  zip
+40    |-           # Zip
+41    |-           )
+42    |-    )
+43    |-  )
+44    |-     (a,  # A
+   36 |+ 
+   37 |+     a,  # A
+45 38 |          b,  # B
+46 39 |          c,  # C
+47    |-   )      )
+48    |-   ),
+   40 |+   
+49 41 | )
+50 42 | 
+51 43 | starmap(
+
+RUF058.py:51:1: RUF058 [*] `itertools.starmap` called on `zip` iterable
+   |
+49 |   )
+50 |
+51 | / starmap(
+52 | |     func \
+53 | |     , \
+54 | |     zip \
+55 | |         (
+56 | |             a,\
+57 | |             b,\
+58 | |             c,\
+59 | |         )
+60 | | )
+   | |_^ RUF058
+   |
+   = help: Use `map` instead
+
+ℹ Safe fix
+48 48 |    ),
+49 49 | )
+50 50 | 
+51    |-starmap(
+   51 |+map(
+52 52 |     func \
+53 53 |     , \
+54    |-    zip \
+55    |-        (
+   54 |+     \
+   55 |+        
+56 56 |             a,\
+57 57 |             b,\
+58 58 |             c,\
+59    |-        )
+   59 |+        
+60 60 | )
+61 61 | 
+62 62 |

crates/ruff_python_ast/src/nodes.rs

@@ -3892,6 +3892,18 @@ impl Arguments {
         let keywords = self.keywords.iter().map(ArgOrKeyword::Keyword);
         args.merge_by(keywords, |left, right| left.start() < right.start())
     }
+
+    pub fn inner_range(&self) -> TextRange {
+        TextRange::new(self.l_paren_range().end(), self.r_paren_range().start())
+    }
+
+    pub fn l_paren_range(&self) -> TextRange {
+        TextRange::at(self.start(), '('.text_len())
+    }
+
+    pub fn r_paren_range(&self) -> TextRange {
+        TextRange::new(self.end() - ')'.text_len(), self.end())
+    }
 }
 
 /// An AST node used to represent a sequence of type parameters.

crates/ruff_text_size/src/range.rs

@@ -44,6 +44,7 @@ impl TextRange {
     /// assert_eq!(range.len(), end - start);
     /// ```
     #[inline]
+    #[track_caller]
     pub const fn new(start: TextSize, end: TextSize) -> TextRange {
         assert!(start.raw <= end.raw);
         TextRange { start, end }

ruff.schema.json

@@ -3903,6 +3903,7 @@
         "RUF055",
         "RUF056",
         "RUF057",
+        "RUF058",
         "RUF1",
         "RUF10",
         "RUF100",