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

crates/ruff_python_formatter/src/expression/expr_list_comp.rs

@@ -1,20 +1,41 @@
 use crate::comments::Comments;
 use crate::expression::parentheses::{
-    default_expression_needs_parentheses, NeedsParentheses, Parentheses, Parenthesize,
+    default_expression_needs_parentheses, parenthesized, NeedsParentheses, Parentheses,
+    Parenthesize,
 };
-use crate::{not_yet_implemented_custom_text, FormatNodeRule, PyFormatter};
-use ruff_formatter::{write, Buffer, FormatResult};
+use crate::prelude::*;
+use crate::AsFormat;
+use crate::{FormatNodeRule, PyFormatter};
+use ruff_formatter::{format_args, write, Buffer, FormatResult};
 use rustpython_parser::ast::ExprListComp;
 
 #[derive(Default)]
 pub struct FormatExprListComp;
 
 impl FormatNodeRule<ExprListComp> for FormatExprListComp {
-    fn fmt_fields(&self, _item: &ExprListComp, f: &mut PyFormatter) -> FormatResult<()> {
+    fn fmt_fields(&self, item: &ExprListComp, f: &mut PyFormatter) -> FormatResult<()> {
+        let ExprListComp {
+            range: _,
+            elt,
+            generators,
+        } = item;
+
+        let joined = format_with(|f| {
+            f.join_with(soft_line_break_or_space())
+                .entries(generators.iter().formatted())
+                .finish()
+        });
+
         write!(
             f,
-            [not_yet_implemented_custom_text(
-                "[NOT_YET_IMPLEMENTED_generator_key for NOT_YET_IMPLEMENTED_generator_key in []]"
+            [parenthesized(
+                "[",
+                &format_args!(
+                    group(&elt.format()),
+                    soft_line_break_or_space(),
+                    group(&joined)
+                ),
+                "]"
             )]
         )
     }

crates/ruff_python_formatter/src/other/comprehension.rs

@@ -1,12 +1,87 @@
-use crate::{not_yet_implemented, FormatNodeRule, PyFormatter};
-use ruff_formatter::{write, Buffer, FormatResult};
-use rustpython_parser::ast::Comprehension;
+use crate::comments::{leading_comments, trailing_comments};
+use crate::prelude::*;
+use crate::AsFormat;
+use crate::{FormatNodeRule, PyFormatter};
+use ruff_formatter::{format_args, write, Buffer, FormatResult};
+use rustpython_parser::ast::{Comprehension, Ranged};
 
 #[derive(Default)]
 pub struct FormatComprehension;
 
 impl FormatNodeRule<Comprehension> for FormatComprehension {
     fn fmt_fields(&self, item: &Comprehension, f: &mut PyFormatter) -> FormatResult<()> {
-        write!(f, [not_yet_implemented(item)])
+        let Comprehension {
+            range: _,
+            target,
+            iter,
+            ifs,
+            is_async,
+        } = item;
+
+        let comments = f.context().comments().clone();
+
+        if *is_async {
+            write!(f, [text("async"), space()])?;
+        }
+
+        let dangling_item_comments = comments.dangling_comments(item);
+
+        let (before_target_comments, before_in_comments) = dangling_item_comments.split_at(
+            dangling_item_comments
+                .partition_point(|comment| comment.slice().end() < target.range().start()),
+        );
+
+        let trailing_in_comments = comments.dangling_comments(iter);
+        write!(
+            f,
+            [
+                text("for"),
+                trailing_comments(before_target_comments),
+                group(&format_args!(
+                    soft_line_break_or_space(),
+                    target.format(),
+                    soft_line_break_or_space(),
+                    leading_comments(before_in_comments),
+                    text("in"),
+                    trailing_comments(trailing_in_comments),
+                    soft_line_break_or_space(),
+                    iter.format(),
+                )),
+            ]
+        )?;
+        if !ifs.is_empty() {
+            let joined = format_with(|f| {
+                let mut joiner = f.join_with(soft_line_break_or_space());
+                for if_case in ifs {
+                    let dangling_if_comments = comments.dangling_comments(if_case);
+
+                    let (own_line_if_comments, end_of_line_if_comments) = dangling_if_comments
+                        .split_at(
+                            dangling_if_comments
+                                .partition_point(|comment| comment.line_position().is_own_line()),
+                        );
+                    joiner.entry(&group(&format_args!(
+                        leading_comments(own_line_if_comments),
+                        text("if"),
+                        trailing_comments(end_of_line_if_comments),
+                        soft_line_break_or_space(),
+                        if_case.format(),
+                    )));
+                }
+                joiner.finish()
+            });
+
+            write!(f, [soft_line_break_or_space(), group(&joined)])?;
+        }
+        Ok(())
+    }
+
+    fn fmt_dangling_comments(
+        &self,
+        _node: &Comprehension,
+        _f: &mut PyFormatter,
+    ) -> FormatResult<()> {
+        // dangling comments are formatted as part of fmt_fields
+        Ok(())
     }
 }