format ExprListComp (#5600)

Co-authored-by: Micha Reiser <micha@reiser.io>

crates/ruff_python_formatter/src/comments/placement.rs

@@ -38,6 +38,7 @@ pub(super) fn place_comment<'a>(
         handle_slice_comments,
         handle_attribute_comment,
         handle_expr_if_comment,
+        handle_comprehension_comment,
         handle_trailing_expression_starred_star_end_of_line_comment,
     ];
     for handler in HANDLERS {
@@ -1244,6 +1245,137 @@ fn find_only_token_in_range(range: TextRange, locator: &Locator, token_kind: Tok
     token
 }
 
+// Handle comments inside comprehensions, e.g.
+//
+// ```python
+// [
+//      a
+//      for  # dangling on the comprehension
+//      b
+//      # dangling on the comprehension
+//      in  # dangling on comprehension.iter
+//      # leading on the iter
+//      c
+//      # dangling on comprehension.if.n
+//      if  # dangling on comprehension.if.n
+//      d
+// ]
+// ```
+fn handle_comprehension_comment<'a>(
+    comment: DecoratedComment<'a>,
+    locator: &Locator,
+) -> CommentPlacement<'a> {
+    let AnyNodeRef::Comprehension(comprehension) = comment.enclosing_node() else {
+        return CommentPlacement::Default(comment);
+    };
+    let is_own_line = comment.line_position().is_own_line();
+
+    // Comments between the `for` and target
+    // ```python
+    // [
+    //      a
+    //      for  # attache as dangling on the comprehension
+    //      b in c
+    //  ]
+    // ```
+    if comment.slice().end() < comprehension.target.range().start() {
+        return if is_own_line {
+            // own line comments are correctly assigned as leading the target
+            CommentPlacement::Default(comment)
+        } else {
+            // after the `for`
+            CommentPlacement::dangling(comment.enclosing_node(), comment)
+        };
+    }
+
+    let in_token = find_only_token_in_range(
+        TextRange::new(
+            comprehension.target.range().end(),
+            comprehension.iter.range().start(),
+        ),
+        locator,
+        TokenKind::In,
+    );
+
+    // Comments between the target and the `in`
+    // ```python
+    // [
+    //      a for b
+    //      # attach as dangling on the target
+    //      # (to be rendered as leading on the "in")
+    //      in c
+    //  ]
+    // ```
+    if comment.slice().start() < in_token.start() {
+        // attach as dangling comments on the target
+        // (to be rendered as leading on the "in")
+        return if is_own_line {
+            CommentPlacement::dangling(comment.enclosing_node(), comment)
+        } else {
+            // correctly trailing on the target
+            CommentPlacement::Default(comment)
+        };
+    }
+
+    // Comments between the `in` and the iter
+    // ```python
+    // [
+    //      a for b
+    //      in  #  attach as dangling on the iter
+    //      c
+    //  ]
+    // ```
+    if comment.slice().start() < comprehension.iter.range().start() {
+        return if is_own_line {
+            CommentPlacement::Default(comment)
+        } else {
+            // after the `in` but same line, turn into trailing on the `in` token
+            CommentPlacement::dangling((&comprehension.iter).into(), comment)
+        };
+    }
+
+    let mut last_end = comprehension.iter.range().end();
+
+    for if_node in &comprehension.ifs {
+        // ```python
+        // [
+        //     a
+        //     for
+        //     c
+        //     in
+        //     e
+        //     # above if   <-- find these own-line between previous and `if` token
+        //     if  # if     <-- find these end-of-line between `if` and if node (`f`)
+        //     # above f    <-- already correctly assigned as leading `f`
+        //     f  # f       <-- already correctly assigned as trailing `f`
+        //     # above if2
+        //     if  # if2
+        //     # above g
+        //     g  # g
+        // ]
+        // ```
+        let if_token = find_only_token_in_range(
+            TextRange::new(last_end, if_node.range().start()),
+            locator,
+            TokenKind::If,
+        );
+        if is_own_line {
+            if last_end < comment.slice().start() && comment.slice().start() < if_token.start() {
+                return CommentPlacement::dangling((if_node).into(), comment);
+            }
+        } else {
+            if if_token.start() < comment.slice().start()
+                && comment.slice().start() < if_node.range().start()
+            {
+                return CommentPlacement::dangling((if_node).into(), comment);
+            }
+        }
+        last_end = if_node.range().end();
+    }
+
+    CommentPlacement::Default(comment)
+}
+
 /// Returns `true` if `right` is `Some` and `left` and `right` are referentially equal.
 fn are_same_optional<'a, T>(left: AnyNodeRef, right: Option<T>) -> bool
 where