Skip to content

Commit

Permalink
Avoid suggesting starmap when arguments are used outside call
Browse files Browse the repository at this point in the history
  • Loading branch information
charliermarsh committed Jun 10, 2024
1 parent 7509a48 commit 6532e6e
Show file tree
Hide file tree
Showing 2 changed files with 26 additions and 7 deletions.
4 changes: 4 additions & 0 deletions crates/ruff_linter/resources/test/fixtures/refurb/FURB140.py
Original file line number Diff line number Diff line change
Expand Up @@ -53,3 +53,7 @@ def zipped():
[print(x, *y) for x, y in zipped()]

[print(*x, *y) for x, y in zipped()]

[" ".join(x)(x, y) for x, y in zipped()]

[" ".join(x)(*x) for x in zipped()]
29 changes: 22 additions & 7 deletions crates/ruff_linter/src/rules/refurb/rules/reimplemented_starmap.rs
Original file line number Diff line number Diff line change
Expand Up @@ -2,6 +2,7 @@ use anyhow::{bail, Result};
use ruff_diagnostics::{Diagnostic, Edit, Fix, FixAvailability, Violation};
use ruff_macros::{derive_message_formats, violation};
use ruff_python_ast::comparable::ComparableExpr;
use ruff_python_ast::helpers::any_over_expr;
use ruff_python_ast::{self as ast, Expr};
use ruff_text_size::{Ranged, TextRange};

Expand Down Expand Up @@ -122,6 +123,13 @@ pub(crate) fn reimplemented_starmap(checker: &mut Checker, target: &StarmapCandi
if ComparableExpr::from(value.as_ref()) != ComparableExpr::from(name) {
return;
}

// If the argument is used outside the function call, we can't replace it.
if any_over_expr(&func, &|expr| {
expr.as_name_expr().is_some_and(|expr| expr.id == name.id)
}) {
return;
}
}
// Ex) `f(x, y, z) for x, y, z in iter`
ComprehensionTarget::Tuple(tuple) => {
Expand All @@ -131,23 +139,30 @@ pub(crate) fn reimplemented_starmap(checker: &mut Checker, target: &StarmapCandi
{
return;
}

// If any of the members are used outside the function call, we can't replace it.
if any_over_expr(&func, &|expr| {
expr.as_name_expr().is_some_and(|expr| {
tuple
.elts
.iter()
.any(|elt| elt.as_name_expr().is_some_and(|elt| elt.id == expr.id))
})
}) {
return;
}
}
}

let mut diagnostic = Diagnostic::new(ReimplementedStarmap, target.range());
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.
// Import `starmap` from `itertools`.
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.
//
// 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.
Expand Down

0 comments on commit 6532e6e

Please sign in to comment.