Avoid suggesting starmap when arguments are used outside call (#11830)

## Summary

Closes https://github.com/astral-sh/ruff/issues/11810.
This commit is contained in:
Charlie Marsh 2024-06-10 14:10:06 -07:00 committed by GitHub
parent 0d06900cec
commit 08b548626a
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
2 changed files with 24 additions and 7 deletions

View file

@ -53,3 +53,7 @@ from itertools import starmap as sm
[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()]

View file

@ -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};
@ -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) => {
@ -131,23 +139,28 @@ 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| {
tuple
.elts
.iter()
.any(|elt| ComparableExpr::from(expr) == ComparableExpr::from(elt))
}) {
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.