Skip to content

Commit

Permalink
Preserve whitespace around ListComp brackets in C419 (#4099)
Browse files Browse the repository at this point in the history
  • Loading branch information
@dhruvmanila
dhruvmanila committed May 9, 2023
1 parent 83536cf commit 085fd37
Show file tree
Hide file tree
Showing 3 changed files with 230 additions and 8 deletions.
20 changes: 20 additions & 0 deletions crates/ruff/resources/test/fixtures/flake8_comprehensions/C419.py
View file
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
View file
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
// 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
63 changes: 63 additions & 0 deletions ...ke8_comprehensions/snapshots/ruff__rules__flake8_comprehensions__tests__C419_C419.py.snap
View file
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 | )


0 comments on commit 085fd37

Please sign in to comment.