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

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
+)

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,
+    Arg, AssignEqual, AssignTargetExpression, Call, Codegen, CodegenState, Comment, CompFor, Dict,
-    DictElement, Element, Expr, Expression, GeneratorExp, LeftCurlyBrace, LeftParen,
+    DictComp, DictElement, Element, EmptyLine, Expr, Expression, GeneratorExp, LeftCurlyBrace,
-    LeftSquareBracket, List, ListComp, Name, ParenthesizableWhitespace, RightCurlyBrace,
+    LeftParen, LeftSquareBracket, List, ListComp, Name, ParenthesizableWhitespace,
-    RightParen, RightSquareBracket, Set, SetComp, SimpleString, SimpleWhitespace, Tuple,
+    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 {
+    let whitespace_after_arg = match &call.args[0].comma {
-        call.args[0].whitespace_after_arg = comma.whitespace_after.clone();
+        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(),

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 | )
+