format ExprListComp (#5600)

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

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(())
     }
 }