From e129f77bcf8a572a5ba544dfcbd9a26f28c598a1 Mon Sep 17 00:00:00 2001 From: Tom Kuson Date: Tue, 3 Oct 2023 02:38:05 +0100 Subject: [PATCH] Extend `reimplemented-starmap` (`FURB140`) to catch calls with a single and starred argument (#7768) --- .../resources/test/fixtures/refurb/FURB140.py | 12 ++++ .../refurb/rules/reimplemented_starmap.rs | 57 ++++++++++++---- ...es__refurb__tests__FURB140_FURB140.py.snap | 68 ++++++++++++++++++- crates/ruff_python_ast/src/comparable.rs | 14 ++-- 4 files changed, 128 insertions(+), 23 deletions(-) diff --git a/crates/ruff_linter/resources/test/fixtures/refurb/FURB140.py b/crates/ruff_linter/resources/test/fixtures/refurb/FURB140.py index fd0bc8461316b..e4de0667c76ea 100644 --- a/crates/ruff_linter/resources/test/fixtures/refurb/FURB140.py +++ b/crates/ruff_linter/resources/test/fixtures/refurb/FURB140.py @@ -24,6 +24,12 @@ def zipped(): # FURB140 {print(x, y) for x, y in zipped()} +# FURB140 (check it still flags starred arguments). +# See https://github.com/astral-sh/ruff/issues/7636 +[foo(*t) for t in [(85, 60), (100, 80)]] +(foo(*t) for t in [(85, 60), (100, 80)]) +{foo(*t) for t in [(85, 60), (100, 80)]} + # Non-errors. [print(x, int) for x, _ in zipped()] @@ -41,3 +47,9 @@ def zipped(): [print() for x, y in zipped()] [print(x, end=y) for x, y in zipped()] + +[print(*x, y) for x, y in zipped()] + +[print(x, *y) for x, y in zipped()] + +[print(*x, *y) for x, y in zipped()] diff --git a/crates/ruff_linter/src/rules/refurb/rules/reimplemented_starmap.rs b/crates/ruff_linter/src/rules/refurb/rules/reimplemented_starmap.rs index a8ce08f46f839..892895337946d 100644 --- a/crates/ruff_linter/src/rules/refurb/rules/reimplemented_starmap.rs +++ b/crates/ruff_linter/src/rules/refurb/rules/reimplemented_starmap.rs @@ -82,7 +82,7 @@ pub(crate) fn reimplemented_starmap(checker: &mut Checker, target: &StarmapCandi // ``` // // `x, y, z, ...` are what we call `elts` for short. - let Some((elts, iter)) = match_comprehension(comprehension) else { + let Some(value) = match_comprehension_target(comprehension) else { return; }; @@ -99,13 +99,30 @@ pub(crate) fn reimplemented_starmap(checker: &mut Checker, target: &StarmapCandi 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; + match value { + // Ex) `f(*x) for x in iter` + ComprehensionTarget::Name(name) => { + let [arg] = args else { + return; + }; + + let Expr::Starred(ast::ExprStarred { value, .. }) = arg else { + return; + }; + + if ComparableExpr::from(value.as_ref()) != ComparableExpr::from(name) { + return; + } + } + // Ex) `f(x, y, z) for x, y, z in iter` + ComprehensionTarget::Tuple(tuple) => { + if tuple.elts.len() != args.len() + || !std::iter::zip(&tuple.elts, args) + .all(|(x, y)| ComparableExpr::from(x) == ComparableExpr::from(y)) + { + return; + } + } } let mut diagnostic = Diagnostic::new(ReimplementedStarmap, target.range()); @@ -127,7 +144,7 @@ pub(crate) fn reimplemented_starmap(checker: &mut Checker, target: &StarmapCandi // - 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.try_make_suggestion(starmap_name, &comprehension.iter, func, checker)?, target.range(), ); Ok(Fix::suggested_edits(import_edit, [main_edit])) @@ -252,7 +269,7 @@ fn try_construct_call( // 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")) + bail!("Can't use built-in `{builtin}` constructor") } // In general, we replace: @@ -306,14 +323,24 @@ fn wrap_with_call_to(call: ast::ExprCall, func_name: &str) -> ast::ExprCall { } } -/// Match that the given comprehension is `(x, y, z, ...) in iter`. -fn match_comprehension(comprehension: &ast::Comprehension) -> Option<(&[Expr], &Expr)> { +#[derive(Debug)] +enum ComprehensionTarget<'a> { + /// E.g., `(x, y, z, ...)` in `(x, y, z, ...) in iter`. + Tuple(&'a ast::ExprTuple), + /// E.g., `x` in `x in iter`. + Name(&'a ast::ExprName), +} + +/// Extract the target from the comprehension (e.g., `(x, y, z)` in `(x, y, z, ...) in iter`). +fn match_comprehension_target(comprehension: &ast::Comprehension) -> Option { if comprehension.is_async || !comprehension.ifs.is_empty() { return None; } - - let ast::ExprTuple { elts, .. } = comprehension.target.as_tuple_expr()?; - Some((elts, &comprehension.iter)) + match &comprehension.target { + Expr::Tuple(tuple) => Some(ComprehensionTarget::Tuple(tuple)), + Expr::Name(name) => Some(ComprehensionTarget::Name(name)), + _ => None, + } } /// Match that the given expression is `func(x, y, z, ...)`. diff --git a/crates/ruff_linter/src/rules/refurb/snapshots/ruff_linter__rules__refurb__tests__FURB140_FURB140.py.snap b/crates/ruff_linter/src/rules/refurb/snapshots/ruff_linter__rules__refurb__tests__FURB140_FURB140.py.snap index 1e7691809d56d..89541187db04b 100644 --- a/crates/ruff_linter/src/rules/refurb/snapshots/ruff_linter__rules__refurb__tests__FURB140_FURB140.py.snap +++ b/crates/ruff_linter/src/rules/refurb/snapshots/ruff_linter__rules__refurb__tests__FURB140_FURB140.py.snap @@ -119,7 +119,7 @@ FURB140.py:25:1: FURB140 [*] Use `itertools.starmap` instead of the generator 25 | {print(x, y) for x, y in zipped()} | ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ FURB140 26 | -27 | # Non-errors. +27 | # FURB140 (check it still flags starred arguments). | = help: Replace with `itertools.starmap` @@ -130,7 +130,69 @@ FURB140.py:25:1: FURB140 [*] Use `itertools.starmap` instead of the generator 25 |-{print(x, y) for x, y in zipped()} 25 |+set(sm(print, zipped())) 26 26 | -27 27 | # Non-errors. -28 28 | +27 27 | # FURB140 (check it still flags starred arguments). +28 28 | # See https://github.com/astral-sh/ruff/issues/7636 + +FURB140.py:29:1: FURB140 [*] Use `itertools.starmap` instead of the generator + | +27 | # FURB140 (check it still flags starred arguments). +28 | # See https://github.com/astral-sh/ruff/issues/7636 +29 | [foo(*t) for t in [(85, 60), (100, 80)]] + | ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ FURB140 +30 | (foo(*t) for t in [(85, 60), (100, 80)]) +31 | {foo(*t) for t in [(85, 60), (100, 80)]} + | + = help: Replace with `itertools.starmap` + +ℹ Suggested fix +26 26 | +27 27 | # FURB140 (check it still flags starred arguments). +28 28 | # See https://github.com/astral-sh/ruff/issues/7636 +29 |-[foo(*t) for t in [(85, 60), (100, 80)]] + 29 |+list(sm(foo, [(85, 60), (100, 80)])) +30 30 | (foo(*t) for t in [(85, 60), (100, 80)]) +31 31 | {foo(*t) for t in [(85, 60), (100, 80)]} +32 32 | + +FURB140.py:30:1: FURB140 [*] Use `itertools.starmap` instead of the generator + | +28 | # See https://github.com/astral-sh/ruff/issues/7636 +29 | [foo(*t) for t in [(85, 60), (100, 80)]] +30 | (foo(*t) for t in [(85, 60), (100, 80)]) + | ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ FURB140 +31 | {foo(*t) for t in [(85, 60), (100, 80)]} + | + = help: Replace with `itertools.starmap` + +ℹ Suggested fix +27 27 | # FURB140 (check it still flags starred arguments). +28 28 | # See https://github.com/astral-sh/ruff/issues/7636 +29 29 | [foo(*t) for t in [(85, 60), (100, 80)]] +30 |-(foo(*t) for t in [(85, 60), (100, 80)]) + 30 |+sm(foo, [(85, 60), (100, 80)]) +31 31 | {foo(*t) for t in [(85, 60), (100, 80)]} +32 32 | +33 33 | # Non-errors. + +FURB140.py:31:1: FURB140 [*] Use `itertools.starmap` instead of the generator + | +29 | [foo(*t) for t in [(85, 60), (100, 80)]] +30 | (foo(*t) for t in [(85, 60), (100, 80)]) +31 | {foo(*t) for t in [(85, 60), (100, 80)]} + | ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ FURB140 +32 | +33 | # Non-errors. + | + = help: Replace with `itertools.starmap` + +ℹ Suggested fix +28 28 | # See https://github.com/astral-sh/ruff/issues/7636 +29 29 | [foo(*t) for t in [(85, 60), (100, 80)]] +30 30 | (foo(*t) for t in [(85, 60), (100, 80)]) +31 |-{foo(*t) for t in [(85, 60), (100, 80)]} + 31 |+set(sm(foo, [(85, 60), (100, 80)])) +32 32 | +33 33 | # Non-errors. +34 34 | diff --git a/crates/ruff_python_ast/src/comparable.rs b/crates/ruff_python_ast/src/comparable.rs index 746c7a5418c8b..b3599402f2f60 100644 --- a/crates/ruff_python_ast/src/comparable.rs +++ b/crates/ruff_python_ast/src/comparable.rs @@ -927,11 +927,7 @@ impl<'a> From<&'a ast::Expr> for ComparableExpr<'a> { }) => Self::Starred(ExprStarred { value: value.into(), }), - ast::Expr::Name(ast::ExprName { - id, - ctx: _, - range: _, - }) => Self::Name(ExprName { id: id.as_str() }), + ast::Expr::Name(name) => name.into(), ast::Expr::List(ast::ExprList { elts, ctx: _, @@ -968,6 +964,14 @@ impl<'a> From<&'a ast::Expr> for ComparableExpr<'a> { } } +impl<'a> From<&'a ast::ExprName> for ComparableExpr<'a> { + fn from(expr: &'a ast::ExprName) -> Self { + Self::Name(ExprName { + id: expr.id.as_str(), + }) + } +} + #[derive(Debug, PartialEq, Eq, Hash)] pub struct StmtFunctionDef<'a> { is_async: bool,