Improve comprehension line break beheavior

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

@@ -30,3 +30,16 @@
     # above g
     g  # g
 ]
+
+[
+    aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa + bbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbb + [dddddddddddddddddd, eeeeeeeeeeeeeeeeeee]
+    for
+    ccccccccccccccccccccccccccccccccccccccc,
+    ddddddddddddddddddd, [eeeeeeeeeeeeeeeeeeeeee, fffffffffffffffffffffffff]
+    in
+    eeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeffffffffffffffffffffffffggggggggggggggggggggghhhhhhhhhhhhhhothermoreeand_even_moreddddddddddddddddddddd
+    if
+    fffffffffffffffffffffffffffffffffffffffffff < gggggggggggggggggggggggggggggggggggggggggggggg < hhhhhhhhhhhhhhhhhhhhhhhhhh
+    if
+    gggggggggggggggggggggggggggggggggggggggggggg
+]

crates/ruff_python_formatter/src/expression/expr_compare.rs

@@ -33,36 +33,40 @@ impl FormatNodeRule<ExprCompare> for FormatExprCompare {
 
         let comments = f.context().comments().clone();
 
-        write!(f, [in_parentheses_only_group(&left.format())])?;
+        let inner = format_with(|f| {
+            write!(f, [in_parentheses_only_group(&left.format())])?;
+
+            assert_eq!(comparators.len(), ops.len());
+
+            for (operator, comparator) in ops.iter().zip(comparators) {
+                let leading_comparator_comments = comments.leading_comments(comparator);
+                if leading_comparator_comments.is_empty() {
+                    write!(f, [soft_line_break_or_space()])?;
+                } else {
+                    // Format the expressions leading comments **before** the operator
+                    write!(
+                        f,
+                        [
+                            hard_line_break(),
+                            leading_comments(leading_comparator_comments)
+                        ]
+                    )?;
+                }
 
-        assert_eq!(comparators.len(), ops.len());
-
-        for (operator, comparator) in ops.iter().zip(comparators) {
-            let leading_comparator_comments = comments.leading_comments(comparator);
-            if leading_comparator_comments.is_empty() {
-                write!(f, [soft_line_break_or_space()])?;
-            } else {
-                // Format the expressions leading comments **before** the operator
                 write!(
                     f,
                     [
-                        hard_line_break(),
-                        leading_comments(leading_comparator_comments)
+                        operator.format(),
+                        space(),
+                        in_parentheses_only_group(&comparator.format())
                     ]
                 )?;
             }
 
-            write!(
-                f,
-                [
-                    operator.format(),
-                    space(),
-                    in_parentheses_only_group(&comparator.format())
-                ]
-            )?;
-        }
-
-        Ok(())
+            Ok(())
+        });
+
+        in_parentheses_only_group(&inner).fmt(f)
     }
 }
 

crates/ruff_python_formatter/src/lib.rs

@@ -280,9 +280,14 @@ if True:
     #[test]
     fn quick_test() {
         let src = r#"
-def foo() -> tuple[int, int, int,]:
-    return 2
-
+lcomp3 = [
+    # This one is actually too long to fit in a single line.
+    element.split("\n", 1)[0]
+    # yup
+    for element in collection.select_elements()
+    # right
+    if element is not None
+]
 "#;
         // Tokenize once
         let mut tokens = Vec::new();
@@ -310,10 +315,10 @@ def foo() -> tuple[int, int, int,]:
         // Uncomment the `dbg` to print the IR.
         // Use `dbg_write!(f, []) instead of `write!(f, [])` in your formatting code to print some IR
         // inside of a `Format` implementation
-        // use ruff_formatter::FormatContext;
-        // dbg!(formatted
-        //     .document()
-        //     .display(formatted.context().source_code()));
+        use ruff_formatter::FormatContext;
+        dbg!(formatted
+            .document()
+            .display(formatted.context().source_code()));
         //
         // dbg!(formatted
         //     .context()

crates/ruff_python_formatter/src/other/comprehension.rs

@@ -3,7 +3,7 @@ use crate::prelude::*;
 use crate::AsFormat;
 use crate::{FormatNodeRule, PyFormatter};
 use ruff_formatter::{format_args, write, Buffer, FormatResult};
-use rustpython_parser::ast::{Comprehension, Ranged};
+use rustpython_parser::ast::{Comprehension, Expr, Ranged};
 
 #[derive(Default)]
 pub struct FormatComprehension;
@@ -18,33 +18,52 @@ impl FormatNodeRule<Comprehension> for FormatComprehension {
             is_async,
         } = item;
 
-        let comments = f.context().comments().clone();
-
         if *is_async {
             write!(f, [text("async"), space()])?;
         }
 
+        let comments = f.context().comments().clone();
         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);
+
+        struct Spacer<'a>(&'a Expr);
+
+        impl Format<PyFormatContext<'_>> for Spacer<'_> {
+            fn fmt(&self, f: &mut PyFormatter) -> FormatResult<()> {
+                if f.context().comments().has_leading_comments(self.0) {
+                    soft_line_break_or_space().fmt(f)
+                } else {
+                    space().fmt(f)
+                }
+            }
+        }
+
+        let in_spacer = format_with(|f| {
+            if before_in_comments.is_empty() {
+                space().fmt(f)
+            } else {
+                soft_line_break_or_space().fmt(f)
+            }
+        });
+
         write!(
             f,
             [
                 text("for"),
                 trailing_comments(before_target_comments),
                 group(&format_args!(
-                    soft_line_break_or_space(),
+                    Spacer(target),
                     target.format(),
-                    soft_line_break_or_space(),
+                    in_spacer,
                     leading_comments(before_in_comments),
                     text("in"),
                     trailing_comments(trailing_in_comments),
-                    soft_line_break_or_space(),
+                    Spacer(iter),
                     iter.format(),
                 )),
             ]
@@ -64,7 +83,7 @@ impl FormatNodeRule<Comprehension> for FormatComprehension {
                         leading_comments(own_line_if_comments),
                         text("if"),
                         trailing_comments(end_of_line_if_comments),
-                        soft_line_break_or_space(),
+                        Spacer(if_case),
                         if_case.format(),
                     )));
                 }

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,27 @@
+@@ -2,14 +2,11 @@
  
  
  def f():
@@ -59,17 +59,7 @@ def make_arange(n):
  
  
  async def func():
-     if test:
-         out_batched = [
-             i
--            async for i in aitertools._async_map(
--                self.async_inc, arange(8), batch_size=3
--            )
-+            async for
-+            i
-+            in
-+            aitertools._async_map(self.async_inc, arange(8), batch_size=3)
-         ]
+@@ -23,8 +20,8 @@
  
  
  def awaited_generator_value(n):
@@ -100,10 +90,9 @@ async def func():
     if test:
         out_batched = [
             i
-            async for
-            i
-            in
-            aitertools._async_map(self.async_inc, arange(8), batch_size=3)
+            async for i in aitertools._async_map(
+                self.async_inc, arange(8), batch_size=3
+            )
         ]
 
 

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

@@ -224,43 +224,18 @@ instruction()#comment with bad spacing
      if (
          self._proc is not None
          # has the child process finished?
-@@ -115,7 +123,12 @@
+@@ -115,7 +123,9 @@
          arg3=True,
      )
      lcomp = [
 -        element for element in collection if element is not None  # yup  # yup  # right
 +        element  # yup
-+        for
-+        element
-+        in
-+        collection  # yup
++        for element in collection  # yup
 +        if element is not None  # right
      ]
      lcomp2 = [
          # hello
-@@ -123,7 +136,9 @@
-         # yup
-         for element in collection
-         # right
--        if element is not None
-+        if
-+        element
-+        is not None
-     ]
-     lcomp3 = [
-         # This one is actually too long to fit in a single line.
-@@ -131,7 +146,9 @@
-         # yup
-         for element in collection.select_elements()
-         # right
--        if element is not None
-+        if
-+        element
-+        is not None
-     ]
-     while True:
-         if False:
-@@ -143,7 +160,10 @@
+@@ -143,7 +153,10 @@
      # let's return
      return Node(
          syms.simple_stmt,
@@ -272,14 +247,13 @@ instruction()#comment with bad spacing
      )
  
  
-@@ -158,7 +178,11 @@
+@@ -158,7 +171,10 @@
  
  class Test:
      def _init_host(self, parsed) -> None:
 -        if parsed.hostname is None or not parsed.hostname.strip():  # type: ignore
 +        if (
-+            parsed.hostname
-+            is None  # type: ignore
++            parsed.hostname is None  # type: ignore
 +            or not parsed.hostname.strip()
 +        ):
              pass
@@ -416,10 +390,7 @@ short
     )
     lcomp = [
         element  # yup
-        for
-        element
-        in
-        collection  # yup
+        for element in collection  # yup
         if element is not None  # right
     ]
     lcomp2 = [
@@ -428,19 +399,15 @@ short
         # yup
         for element in collection
         # right
-        if
-        element
-        is not None
+        if element is not None
     ]
     lcomp3 = [
         # This one is actually too long to fit in a single line.
         element.split("\n", 1)[0]
         # yup
         for element in collection.select_elements()
         # right
-        if
-        element
-        is not None
+        if element is not None
     ]
     while True:
         if False:
@@ -471,8 +438,7 @@ CONFIG_FILES = (
 class Test:
     def _init_host(self, parsed) -> None:
         if (
-            parsed.hostname
-            is None  # type: ignore
+            parsed.hostname is None  # type: ignore
             or not parsed.hostname.strip()
         ):
             pass