format ExprListComp

crates/ruff_python_formatter/resources/test/fixtures/ruff/statement/list_comp.py

@@ -0,0 +1,25 @@
+[i for i in []]
+
+[i for i in [1,]]
+
+[
+    a  # a
+    for # for
+    c  # c
+    in  # in
+    e  # e
+]
+
+[
+    # before a
+    a  # a
+    # before for
+    for # for
+    # before c
+    c  # c
+    # before in
+    in  # in
+    # before e
+    e
+    # before end
+]

crates/ruff_python_formatter/src/comments/placement.rs

@@ -1,6 +1,6 @@
 use std::cmp::Ordering;
 
-use ruff_text_size::{TextLen, TextRange};
+use ruff_text_size::{TextLen, TextRange, TextSize};
 use rustpython_parser::ast;
 use rustpython_parser::ast::{Expr, ExprIfExp, ExprSlice, Ranged};
 
@@ -40,6 +40,7 @@ pub(super) fn place_comment<'a>(
         handle_slice_comments,
         handle_attribute_comment,
         handle_expr_if_comment,
+        handle_comprehension_comment,
     ];
     for handler in HANDLERS {
         comment = match handler(comment, locator) {
@@ -1233,6 +1234,124 @@ fn find_only_token_str_in_range(range: TextRange, locator: &Locator, token_str:
     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> {
+    fn find_if_token(locator: &Locator, start: TextSize, end: TextSize) -> Option<TextSize> {
+        let mut tokens =
+            SimpleTokenizer::new(locator.contents(), TextRange::new(start, end)).skip_trivia();
+        tokens.next().map(|t| t.start())
+    }
+
+    let AnyNodeRef::Comprehension(comprehension) = comment.enclosing_node() else {
+        return CommentPlacement::Default(comment);
+    };
+    let is_own_line = comment.line_position().is_own_line();
+
+    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 mut tokens = SimpleTokenizer::new(
+        locator.contents(),
+        TextRange::new(
+            comprehension.target.range().end(),
+            comprehension.iter.range().start(),
+        ),
+    )
+    .skip_trivia();
+    let Some(in_token) = tokens.next() else {
+        // Should always have an `in`
+        debug_assert!(false);
+        return CommentPlacement::Default(comment);
+    };
+
+    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)
+        };
+    }
+
+    if comment.slice().start() < comprehension.iter.range().start() {
+        // attach as dangling comments on the iter
+
+        return if is_own_line {
+            // after the `in` and own line: leading on the iter
+            CommentPlacement::leading((&comprehension.iter).into(), comment)
+        } else {
+            // after the `in` but same line, turn into trailin on the `in` token
+            CommentPlacement::dangling((&comprehension.iter).into(), comment)
+        };
+    }
+
+    let mut last_end = comprehension.iter.range().end();
+
+    for if_node in &comprehension.ifs {
+        // [
+        //     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 Some(if_token_start) = find_if_token(locator, last_end, if_node.range().end()) else {
+            // there should always be an `if` here
+            debug_assert!(false);
+            return CommentPlacement::Default(comment);
+        };
+        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();
+    }
+
+    return 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,37 @@
 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 crate::prelude::*;
+use crate::AsFormat;
+use crate::{FormatNodeRule, PyFormatter};
 use ruff_formatter::{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(
+                "[",
+                &group(&self::format_args!(&elt.format(), space(), group(&joined))),
+                "]"
             )]
         )
     }

crates/ruff_python_formatter/src/other/comprehension.rs

@@ -1,12 +1,115 @@
-use crate::{not_yet_implemented, FormatNodeRule, PyFormatter};
+use crate::comments::{leading_comments, trailing_comments, SourceComment};
+use crate::prelude::*;
+use crate::AsFormat;
+use crate::{FormatNodeRule, PyFormatter};
 use ruff_formatter::{write, Buffer, FormatResult};
-use rustpython_parser::ast::Comprehension;
+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();
+        let leading_item_comments = comments.leading_comments(item);
+
+        // TODO: why is this needed?
+        if let Some(leading_item_comment) = leading_item_comments.first() {
+            if leading_item_comment.line_position().is_own_line() {
+                hard_line_break().fmt(f)?;
+            }
+        }
+        leading_comments(leading_item_comments).fmt(f)?;
+
+        if *is_async {
+            write!(f, [text("async"), soft_line_break_or_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(&self::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_ in ifs {
+                    let dangling_if_comments = comments.dangling_comments(if_);
+
+                    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(&self::format_args!(
+                        leading_comments(own_line_if_comments),
+                        text("if"),
+                        trailing_comments(end_of_line_if_comments),
+                        soft_line_break_or_space(),
+                        if_.format(),
+                    )));
+                }
+                joiner.finish()
+            });
+
+            write!(f, [soft_line_break_or_space(), group(&joined)])?;
+        }
+        Ok(())
+    }
+
+    fn fmt_leading_comments(
+        &self,
+        _item: &Comprehension,
+        _f: &mut PyFormatter,
+    ) -> FormatResult<()> {
+        Ok(())
+    }
+}
+
+#[derive(Debug)]
+struct FormatLeadingCommentsSpacing<'a> {
+    comments: &'a [SourceComment],
+}
+
+impl Format<PyFormatContext<'_>> for FormatLeadingCommentsSpacing<'_> {
+    fn fmt(&self, f: &mut PyFormatter) -> FormatResult<()> {
+        if let Some(first) = self.comments.first() {
+            if first.line_position().is_own_line() {
+                // Insert a newline after the colon so the comment ends up on its own line
+                hard_line_break().fmt(f)?;
+            } else {
+                // Insert the two spaces between the colon and the end-of-line comment after the colon
+                write!(f, [space(), space()])?;
+            }
+        }
+        Ok(())
     }
 }

crates/ruff_python_formatter/tests/snapshots/black_compatibility@py_37__python37.py.snap

@@ -42,7 +42,7 @@ def make_arange(n):
 ```diff
 --- Black
 +++ Ruff
-@@ -2,29 +2,21 @@
+@@ -2,29 +2,24 @@
  
  
  def f():
@@ -60,13 +60,14 @@ def make_arange(n):
  
  async def func():
      if test:
--        out_batched = [
+         out_batched = [
 -            i
 -            async for i in aitertools._async_map(
 -                self.async_inc, arange(8), batch_size=3
 -            )
--        ]
-+        out_batched = [NOT_YET_IMPLEMENTED_generator_key for NOT_YET_IMPLEMENTED_generator_key in []]
++            i async
++            for i in aitertools._async_map(self.async_inc, arange(8), batch_size=3)
+         ]
  
  
  def awaited_generator_value(n):
@@ -95,7 +96,10 @@ def g():
 
 async def func():
     if test:
-        out_batched = [NOT_YET_IMPLEMENTED_generator_key for NOT_YET_IMPLEMENTED_generator_key in []]
+        out_batched = [
+            i async
+            for i in aitertools._async_map(self.async_inc, arange(8), batch_size=3)
+        ]
 
 
 def awaited_generator_value(n):

crates/ruff_python_formatter/tests/snapshots/black_compatibility@py_38__pep_572.py.snap

@@ -76,7 +76,7 @@ while x := f(x):
 -y0 = (y1 := f(x))
 -foo(x=(y := f(x)))
 +[NOT_YET_IMPLEMENTED_ExprNamedExpr, y**2, y**3]
-+filtered_data = [NOT_YET_IMPLEMENTED_generator_key for NOT_YET_IMPLEMENTED_generator_key in []]
++filtered_data = [y for x in data if (NOT_YET_IMPLEMENTED_ExprNamedExpr) is None]
 +(NOT_YET_IMPLEMENTED_ExprNamedExpr)
 +y0 = NOT_YET_IMPLEMENTED_ExprNamedExpr
 +foo(x=(NOT_YET_IMPLEMENTED_ExprNamedExpr))
@@ -150,7 +150,7 @@ if (NOT_YET_IMPLEMENTED_ExprNamedExpr) is None:
 if NOT_YET_IMPLEMENTED_ExprNamedExpr:
     pass
 [NOT_YET_IMPLEMENTED_ExprNamedExpr, y**2, y**3]
-filtered_data = [NOT_YET_IMPLEMENTED_generator_key for NOT_YET_IMPLEMENTED_generator_key in []]
+filtered_data = [y for x in data if (NOT_YET_IMPLEMENTED_ExprNamedExpr) is None]
 (NOT_YET_IMPLEMENTED_ExprNamedExpr)
 y0 = NOT_YET_IMPLEMENTED_ExprNamedExpr
 foo(x=(NOT_YET_IMPLEMENTED_ExprNamedExpr))