Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Preserve whitespace around ListComp brackets in C419 #4099

Merged
merged 3 commits into from
May 9, 2023
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
20 changes: 20 additions & 0 deletions crates/ruff/resources/test/fixtures/flake8_comprehensions/C419.py
Original file line number Diff line number Diff line change
Expand Up @@ -17,3 +17,23 @@

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 changes: 147 additions & 8 deletions crates/ruff/src/rules/flake8_comprehensions/fixes.rs
Original file line number Diff line number Diff line change
@@ -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;
Expand Down Expand Up @@ -1156,17 +1157,155 @@ 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
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I love these comments! Thank you so much for taking the time to write them. It makes reviewing (and reading the code in the future) so much easier.

// 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(),
lpar: list_comp.lpar.clone(),
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(),
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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 | )