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..fd0bc8461316b --- /dev/null +++ b/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()] diff --git a/crates/ruff/src/checkers/ast/analyze/expression.rs b/crates/ruff/src/checkers/ast/analyze/expression.rs index b28170a35039c..f490828aa9894 100644 --- a/crates/ruff/src/checkers/ast/analyze/expression.rs +++ b/crates/ruff/src/checkers/ast/analyze/expression.rs @@ -1276,16 +1276,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, @@ -1299,6 +1296,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, @@ -1323,11 +1347,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)); } @@ -1336,6 +1362,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 { diff --git a/crates/ruff/src/codes.rs b/crates/ruff/src/codes.rs index d9d8b7b1b5db5..3770d756598dd 100644 --- a/crates/ruff/src/codes.rs +++ b/crates/ruff/src/codes.rs @@ -918,6 +918,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 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..f8994b4cb995a --- /dev/null +++ b/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 { + 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 { + 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 { + // 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..e07b92ad23293 100644 --- a/crates/ruff_workspace/src/configuration.rs +++ b/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)] diff --git a/ruff.schema.json b/ruff.schema.json index 74f7245fcbbc7..966eab2bf140e 100644 --- a/ruff.schema.json +++ b/ruff.schema.json @@ -2091,6 +2091,7 @@ "FURB131", "FURB132", "FURB14", + "FURB140", "FURB145", "G", "G0",