Support C419 autofixes for set comprehensions (#6744)

Closes https://github.com/astral-sh/ruff/issues/6713.
This commit is contained in:
Charlie Marsh 2023-08-21 19:41:13 -04:00 committed by GitHub
parent 7b14d17e39
commit 7650c6ee45
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
3 changed files with 61 additions and 32 deletions

View File

@ -1045,8 +1045,26 @@ pub(crate) fn fix_unnecessary_comprehension_any_all(
let mut tree = match_expression(module_text)?;
let call = match_call_mut(&mut tree)?;
let Expression::ListComp(list_comp) = &call.args[0].value else {
bail!("Expected Expression::ListComp");
let (whitespace_after, whitespace_before, elt, for_in, lpar, rpar) = match &call.args[0].value {
Expression::ListComp(list_comp) => (
&list_comp.lbracket.whitespace_after,
&list_comp.rbracket.whitespace_before,
&list_comp.elt,
&list_comp.for_in,
&list_comp.lpar,
&list_comp.rpar,
),
Expression::SetComp(set_comp) => (
&set_comp.lbrace.whitespace_after,
&set_comp.rbrace.whitespace_before,
&set_comp.elt,
&set_comp.for_in,
&set_comp.lpar,
&set_comp.rpar,
),
_ => {
bail!("Expected Expression::ListComp | Expression::SetComp");
}
};
let mut new_empty_lines = vec![];
@ -1055,7 +1073,7 @@ pub(crate) fn fix_unnecessary_comprehension_any_all(
first_line,
empty_lines,
..
}) = &list_comp.lbracket.whitespace_after
}) = &whitespace_after
{
// If there's a comment on the line after the opening bracket, we need
// to preserve it. The way we do this is by adding a new empty line
@ -1144,7 +1162,7 @@ pub(crate) fn fix_unnecessary_comprehension_any_all(
..
},
..
}) = &list_comp.rbracket.whitespace_before
}) = &whitespace_before
{
Some(format!("{}{}", whitespace.0, comment.0))
} else {
@ -1152,10 +1170,10 @@ pub(crate) fn fix_unnecessary_comprehension_any_all(
};
call.args[0].value = Expression::GeneratorExp(Box::new(GeneratorExp {
elt: list_comp.elt.clone(),
for_in: list_comp.for_in.clone(),
lpar: list_comp.lpar.clone(),
rpar: list_comp.rpar.clone(),
elt: elt.clone(),
for_in: for_in.clone(),
lpar: lpar.clone(),
rpar: rpar.clone(),
}));
let whitespace_after_arg = match &call.args[0].comma {

View File

@ -69,30 +69,31 @@ pub(crate) fn unnecessary_comprehension_any_all(
let Expr::Name(ast::ExprName { id, .. }) = func else {
return;
};
if (matches!(id.as_str(), "all" | "any")) && args.len() == 1 {
let (Expr::ListComp(ast::ExprListComp { elt, .. })
| Expr::SetComp(ast::ExprSetComp { elt, .. })) = &args[0]
else {
return;
};
if contains_await(elt) {
return;
}
if !checker.semantic().is_builtin(id) {
return;
}
let mut diagnostic = Diagnostic::new(UnnecessaryComprehensionAnyAll, args[0].range());
if checker.patch(diagnostic.kind.rule()) {
diagnostic.try_set_fix(|| {
fixes::fix_unnecessary_comprehension_any_all(
checker.locator(),
checker.stylist(),
expr,
)
});
}
checker.diagnostics.push(diagnostic);
if !matches!(id.as_str(), "all" | "any") {
return;
}
let [arg] = args else {
return;
};
let (Expr::ListComp(ast::ExprListComp { elt, .. })
| Expr::SetComp(ast::ExprSetComp { elt, .. })) = arg
else {
return;
};
if contains_await(elt) {
return;
}
if !checker.semantic().is_builtin(id) {
return;
}
let mut diagnostic = Diagnostic::new(UnnecessaryComprehensionAnyAll, arg.range());
if checker.patch(diagnostic.kind.rule()) {
diagnostic.try_set_fix(|| {
fixes::fix_unnecessary_comprehension_any_all(checker.locator(), checker.stylist(), expr)
});
}
checker.diagnostics.push(diagnostic);
}
/// Return `true` if the [`Expr`] contains an `await` expression.

View File

@ -77,7 +77,7 @@ C419.py:7:5: C419 [*] Unnecessary list comprehension.
9 9 | any({x.id for x in bar})
10 10 |
C419.py:9:5: C419 Unnecessary list comprehension.
C419.py:9:5: C419 [*] Unnecessary list comprehension.
|
7 | [x.id for x in bar], # second comment
8 | ) # third comment
@ -88,6 +88,16 @@ C419.py:9:5: C419 Unnecessary list comprehension.
|
= help: Remove unnecessary list comprehension
Suggested fix
6 6 | all( # first comment
7 7 | [x.id for x in bar], # second comment
8 8 | ) # third comment
9 |-any({x.id for x in bar})
9 |+any(x.id for x in bar)
10 10 |
11 11 | # OK
12 12 | all(x.id for x in bar)
C419.py:24:5: C419 [*] Unnecessary list comprehension.
|
22 | # Special comment handling