diff --git a/crates/ruff/resources/test/fixtures/flake8_comprehensions/C419.py b/crates/ruff/resources/test/fixtures/flake8_comprehensions/C419.py index 74aea42a53..4a9671b1ee 100644 --- a/crates/ruff/resources/test/fixtures/flake8_comprehensions/C419.py +++ b/crates/ruff/resources/test/fixtures/flake8_comprehensions/C419.py @@ -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 +) diff --git a/crates/ruff/src/rules/flake8_comprehensions/fixes.rs b/crates/ruff/src/rules/flake8_comprehensions/fixes.rs index 57a0a78ee9..335174f327 100644 --- a/crates/ruff/src/rules/flake8_comprehensions/fixes.rs +++ b/crates/ruff/src/rules/flake8_comprehensions/fixes.rs @@ -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(), diff --git a/crates/ruff/src/rules/flake8_comprehensions/snapshots/ruff__rules__flake8_comprehensions__tests__C419_C419.py.snap b/crates/ruff/src/rules/flake8_comprehensions/snapshots/ruff__rules__flake8_comprehensions__tests__C419_C419.py.snap index 47df13091c..6b527f29ec 100644 --- a/crates/ruff/src/rules/flake8_comprehensions/snapshots/ruff__rules__flake8_comprehensions__tests__C419_C419.py.snap +++ b/crates/ruff/src/rules/flake8_comprehensions/snapshots/ruff__rules__flake8_comprehensions__tests__C419_C419.py.snap @@ -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 | ) +