diff --git a/crates/ruff/resources/test/fixtures/refurb/FURB140.py b/crates/ruff/resources/test/fixtures/refurb/FURB140.py new file mode 100644 index 0000000000000..82d14f204c9bb --- /dev/null +++ b/crates/ruff/resources/test/fixtures/refurb/FURB140.py @@ -0,0 +1,41 @@ +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, 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()] diff --git a/crates/ruff/src/checkers/ast/analyze/expression.rs b/crates/ruff/src/checkers/ast/analyze/expression.rs index f0b2a05fce4b8..1738378c7881c 100644 --- a/crates/ruff/src/checkers/ast/analyze/expression.rs +++ b/crates/ruff/src/checkers/ast/analyze/expression.rs @@ -1273,16 +1273,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 { - elt, - generators, - range: _, - }) - | Expr::SetComp(ast::ExprSetComp { - elt, - generators, - range: _, - }) => { + Expr::ListComp( + comp @ ast::ExprListComp { + elt, + generators, + range: _, + }, + ) => { if checker.enabled(Rule::UnnecessaryComprehension) { flake8_comprehensions::rules::unnecessary_list_set_comprehension( checker, expr, elt, generators, @@ -1296,6 +1293,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); + } + } + 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); + } } Expr::DictComp(ast::ExprDictComp { key, @@ -1320,11 +1344,13 @@ pub(crate) fn expression(expr: &Expr, checker: &mut Checker) { ruff::rules::static_key_dict_comprehension(checker, key); } } - Expr::GeneratorExp(ast::ExprGeneratorExp { - generators, - elt: _, - range: _, - }) => { + Expr::GeneratorExp( + generator @ ast::ExprGeneratorExp { + generators, + elt: _, + range: _, + }, + ) => { if checker.enabled(Rule::FunctionUsesLoopVariable) { flake8_bugbear::rules::function_uses_loop_variable(checker, &Node::Expr(expr)); } @@ -1333,6 +1359,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); + } } Expr::BoolOp( bool_op @ ast::ExprBoolOp { diff --git a/crates/ruff/src/codes.rs b/crates/ruff/src/codes.rs index df4c8a5f0b8f2..addf22b7d6f1a 100644 --- a/crates/ruff/src/codes.rs +++ b/crates/ruff/src/codes.rs @@ -918,6 +918,8 @@ 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), + #[allow(deprecated)] + (Refurb, "140") => (RuleGroup::Nursery, rules::refurb::rules::ReimplementedStarmap), (Refurb, "145") => (RuleGroup::Preview, rules::refurb::rules::SliceCopy), // flake8-logging diff --git a/crates/ruff/src/rules/refurb/mod.rs b/crates/ruff/src/rules/refurb/mod.rs index a910946595020..27b1c59e9f053 100644 --- a/crates/ruff/src/rules/refurb/mod.rs +++ b/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()); diff --git a/crates/ruff/src/rules/refurb/rules/mod.rs b/crates/ruff/src/rules/refurb/rules/mod.rs index 414a4168f4b50..55a5095e251b2 100644 --- a/crates/ruff/src/rules/refurb/rules/mod.rs +++ b/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; diff --git a/crates/ruff/src/rules/refurb/rules/reimplemented_starmap.rs b/crates/ruff/src/rules/refurb/rules/reimplemented_starmap.rs new file mode 100644 index 0000000000000..af18219707b17 --- /dev/null +++ b/crates/ruff/src/rules/refurb/rules/reimplemented_starmap.rs @@ -0,0 +1,329 @@ +use anyhow::{bail, Result}; +use ruff_diagnostics::{AutofixKind, Diagnostic, Edit, Fix, Violation}; +use ruff_macros::{derive_message_formats, violation}; +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 { + Some(format!("Replace with `itertools.starmap`")) + } +} + +/// A abstract node that can be considered a candidate for replacement with `starmap`. +pub(crate) trait StarmapEquivalent: Ranged { + /// Get generated element. + fn element(&self) -> &Expr; + /// Get generator comprehensions. + fn generators(&self) -> &[ast::Comprehension]; + /// Try to produce a fix suggestion transforming this node into a call to `starmap`. + fn try_make_suggestion( + name: String, + iter: &Expr, + func: &Expr, + checker: &Checker, + ) -> Result; +} + +// FURB140 +pub(crate) fn reimplemented_starmap(checker: &mut Checker, generator: &T) { + // Generator should have exactly one comprehension. + let [comprehension @ ast::Comprehension { .. }] = generator.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(generator.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) + // We intentionally do not use ComparableExpr here because it will compare expression + // contexts and in `elts` it's definitely `Load`, while in `args` it's always `Store`. + // + // For this reason, we compare names directly. + .all(|(x, y)| get_name_id(x) == get_name_id(y)) + { + return; + } + + let mut diagnostic = Diagnostic::new(ReimplementedStarmap, generator.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"), + generator.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( + T::try_make_suggestion(starmap_name, iter, func, checker)?, + generator.range(), + ); + Ok(Fix::suggested_edits(import_edit, [main_edit])) + }); + } + checker.diagnostics.push(diagnostic); +} + +#[inline] +fn get_name_id(expr: &Expr) -> Option<&str> { + Some(&expr.as_name_expr()?.id) +} + +impl StarmapEquivalent for ast::ExprGeneratorExp { + fn element(&self) -> &Expr { + self.elt.as_ref() + } + fn generators(&self) -> &[ast::Comprehension] { + self.generators.as_slice() + } + fn try_make_suggestion( + name: String, + iter: &Expr, + func: &Expr, + checker: &Checker, + ) -> Result { + // 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())) + } +} + +impl StarmapEquivalent for ast::ExprListComp { + fn element(&self) -> &Expr { + self.elt.as_ref() + } + fn generators(&self) -> &[ast::Comprehension] { + self.generators.as_slice() + } + fn try_make_suggestion( + name: String, + iter: &Expr, + func: &Expr, + checker: &Checker, + ) -> Result { + // 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) + } +} + +impl StarmapEquivalent for ast::ExprSetComp { + fn element(&self) -> &Expr { + self.elt.as_ref() + } + fn generators(&self) -> &[ast::Comprehension] { + self.generators.as_slice() + } + fn try_make_suggestion( + name: String, + iter: &Expr, + func: &Expr, + checker: &Checker, + ) -> Result { + // 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 { + // 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)) +} diff --git a/crates/ruff/src/rules/refurb/snapshots/ruff__rules__refurb__tests__FURB140_FURB140.py.snap b/crates/ruff/src/rules/refurb/snapshots/ruff__rules__refurb__tests__FURB140_FURB140.py.snap new file mode 100644 index 0000000000000..8e08c4c5c99c0 --- /dev/null +++ b/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 | + + diff --git a/crates/ruff_workspace/src/configuration.rs b/crates/ruff_workspace/src/configuration.rs index b5778bbf22d17..c362d7e507d95 100644 --- a/crates/ruff_workspace/src/configuration.rs +++ b/crates/ruff_workspace/src/configuration.rs @@ -850,6 +850,7 @@ mod tests { Rule::DeleteFullSlice, Rule::CheckAndRemoveFromSet, Rule::QuadraticListSummation, + Rule::ReimplementedStarmap, ]; const PREVIEW_RULES: &[Rule] = &[ diff --git a/ruff.schema.json b/ruff.schema.json index c09942966a871..b9877a829540c 100644 --- a/ruff.schema.json +++ b/ruff.schema.json @@ -2091,6 +2091,7 @@ "FURB131", "FURB132", "FURB14", + "FURB140", "FURB145", "G", "G0",