Python
/
ruff
mirror of https://github.com/astral-sh/ruff
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

Preserve whitespace around `ListComp` brackets in `C419` (#4099)

This commit is contained in:
Dhruv Manilawala 2023-05-09 12:13:05 +05:30 committed by GitHub
parent 83536cf87b
commit 085fd37209
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
3 changed files with 230 additions and 8 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

20
crates/ruff/resources/test/fixtures/flake8_comprehensions/C419.py vendored
View File

@ -17,3 +17,23 @@ all((x.id for x in bar))
async def f() -> bool:
return all([await use_greeting(greeting) for greeting in await greetings()])
# Special comment handling
any(
[ # lbracket comment
# second line comment
i.bit_count()
# random middle comment
for i in range(5) # rbracket comment
] # rpar comment
# trailing comment
)
# Weird case where the function call, opening bracket, and comment are all
# on the same line.
any([ # lbracket comment
# second line comment
i.bit_count() for i in range(5) # rbracket comment
] # rpar comment
)

155
crates/ruff/src/rules/flake8_comprehensions/fixes.rs
View File

@ -1,10 +1,11 @@
use anyhow::{bail, Result};
use itertools::Itertools;
use libcst_native::{
Arg, AssignEqual, AssignTargetExpression, Call, Codegen, CodegenState, CompFor, Dict, DictComp,
DictElement, Element, Expr, Expression, GeneratorExp, LeftCurlyBrace, LeftParen,
LeftSquareBracket, List, ListComp, Name, ParenthesizableWhitespace, RightCurlyBrace,
RightParen, RightSquareBracket, Set, SetComp, SimpleString, SimpleWhitespace, Tuple,
Arg, AssignEqual, AssignTargetExpression, Call, Codegen, CodegenState, Comment, CompFor, Dict,
DictComp, DictElement, Element, EmptyLine, Expr, Expression, GeneratorExp, LeftCurlyBrace,
LeftParen, LeftSquareBracket, List, ListComp, Name, ParenthesizableWhitespace,
ParenthesizedWhitespace, RightCurlyBrace, RightParen, RightSquareBracket, Set, SetComp,
SimpleString, SimpleWhitespace, TrailingWhitespace, Tuple,
};
use ruff_diagnostics::Edit;
@ -1156,6 +1157,108 @@ pub fn fix_unnecessary_comprehension_any_all(
);
};
let mut new_empty_lines = vec![];
if let ParenthesizableWhitespace::ParenthesizedWhitespace(ParenthesizedWhitespace {
first_line,
empty_lines,
..
}) = &list_comp.lbracket.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
// with the same comment.
//
// Example:
// ```python
// any(
// [ # comment
// ...
// ]
// )
//
// # The above code will be converted to:
// any(
// # comment
// ...
// )
// ```
if let TrailingWhitespace {
comment: Some(comment),
..
} = first_line
{
// The indentation should be same as that of the opening bracket,
// but we don't have that information here. This will be addressed
// before adding these new nodes.
new_empty_lines.push(EmptyLine {
comment: Some(comment.clone()),
..EmptyLine::default()
});
}
if !empty_lines.is_empty() {
new_empty_lines.extend(empty_lines.clone());
}
}
if !new_empty_lines.is_empty() {
call.whitespace_before_args = match &call.whitespace_before_args {
ParenthesizableWhitespace::ParenthesizedWhitespace(ParenthesizedWhitespace {
first_line,
indent,
last_line,
..
}) => {
// Add the indentation of the opening bracket to all the new
// empty lines.
for empty_line in &mut new_empty_lines {
empty_line.whitespace = last_line.clone();
}
ParenthesizableWhitespace::ParenthesizedWhitespace(ParenthesizedWhitespace {
first_line: first_line.clone(),
empty_lines: new_empty_lines,
indent: *indent,
last_line: last_line.clone(),
})
}
// This is a rare case, but it can happen if the opening bracket
// is on the same line as the function call.
//
// Example:
// ```python
// any([
// ...
// ]
// )
// ```
ParenthesizableWhitespace::SimpleWhitespace(whitespace) => {
for empty_line in &mut new_empty_lines {
empty_line.whitespace = whitespace.clone();
}
ParenthesizableWhitespace::ParenthesizedWhitespace(ParenthesizedWhitespace {
empty_lines: new_empty_lines,
..ParenthesizedWhitespace::default()
})
}
}
}
let rbracket_comment =
if let ParenthesizableWhitespace::ParenthesizedWhitespace(ParenthesizedWhitespace {
first_line:
TrailingWhitespace {
whitespace,
comment: Some(comment),
..
},
..
}) = &list_comp.rbracket.whitespace_before
{
Some(format!("{}{}", whitespace.0, comment.0))
} else {
None
};
call.args[0].value = Expression::GeneratorExp(Box::new(GeneratorExp {
elt: list_comp.elt.clone(),
for_in: list_comp.for_in.clone(),
@ -1163,10 +1266,46 @@ pub fn fix_unnecessary_comprehension_any_all(
rpar: list_comp.rpar.clone(),
}));
if let Some(comma) = &call.args[0].comma {
call.args[0].whitespace_after_arg = comma.whitespace_after.clone();
call.args[0].comma = None;
}
let whitespace_after_arg = match &call.args[0].comma {
Some(comma) => {
let whitespace_after_comma = comma.whitespace_after.clone();
call.args[0].comma = None;
whitespace_after_comma
}
_ => call.args[0].whitespace_after_arg.clone(),
};
let new_comment;
call.args[0].whitespace_after_arg = match rbracket_comment {
Some(existing_comment) => {
if let ParenthesizableWhitespace::ParenthesizedWhitespace(ParenthesizedWhitespace {
first_line:
TrailingWhitespace {
whitespace: SimpleWhitespace(whitespace),
comment: Some(Comment(comment)),
..
},
empty_lines,
indent,
last_line,
}) = &whitespace_after_arg
{
new_comment = format!("{existing_comment}{whitespace}{comment}");
ParenthesizableWhitespace::ParenthesizedWhitespace(ParenthesizedWhitespace {
first_line: TrailingWhitespace {
comment: Some(Comment(new_comment.as_str())),
..TrailingWhitespace::default()
},
empty_lines: empty_lines.clone(),
indent: *indent,
last_line: last_line.clone(),
})
} else {
whitespace_after_arg
}
}
_ => whitespace_after_arg,
};
let mut state = CodegenState {
default_newline: &stylist.line_ending(),

63
crates/ruff/src/rules/flake8_comprehensions/snapshots/ruff__rules__flake8_comprehensions__tests__C419_C419.py.snap
View File

@ -88,4 +88,67 @@ C419.py:9:5: C419 [*] Unnecessary list comprehension.
|
= help: Remove unnecessary list comprehension
C419.py:24:5: C419 [*] Unnecessary list comprehension.
|
24 | # Special comment handling
25 | any(
26 | [ # lbracket comment
| _____^
27 | | # second line comment
28 | | i.bit_count()
29 | | # random middle comment
30 | | for i in range(5) # rbracket comment
31 | | ] # rpar comment
| |_____^ C419
32 | # trailing comment
33 | )
|
= help: Remove unnecessary list comprehension
ℹ Suggested fix
21 21 |
22 22 | # Special comment handling
23 23 | any(
24 |- [ # lbracket comment
25 |- # second line comment
26 |- i.bit_count()
24 |+ # lbracket comment
25 |+ # second line comment
26 |+ i.bit_count()
27 27 | # random middle comment
28 |- for i in range(5) # rbracket comment
29 |- ] # rpar comment
28 |+ for i in range(5) # rbracket comment # rpar comment
30 29 | # trailing comment
31 30 | )
32 31 |
C419.py:35:5: C419 [*] Unnecessary list comprehension.
|
35 | # Weird case where the function call, opening bracket, and comment are all
36 | # on the same line.
37 | any([ # lbracket comment
| _____^
38 | | # second line comment
39 | | i.bit_count() for i in range(5) # rbracket comment
40 | | ] # rpar comment
| |_____^ C419
41 | )
|
= help: Remove unnecessary list comprehension
ℹ Suggested fix
32 32 |
33 33 | # Weird case where the function call, opening bracket, and comment are all
34 34 | # on the same line.
35 |-any([ # lbracket comment
36 |- # second line comment
37 |- i.bit_count() for i in range(5) # rbracket comment
38 |- ] # rpar comment
35 |+any(
36 |+# lbracket comment
37 |+# second line comment
38 |+i.bit_count() for i in range(5) # rbracket comment # rpar comment
39 39 | )