Format SetComp (astral-sh#5774)

## Summary Format `SetComp` like `ListComp`. ## Test Plan Derived from `ListComp`'s fixture.
evanrittenhouse · Jul 19, 2023 · 1cc2c87 · 1cc2c87
1 parent 8a6afe2
commit 1cc2c87
Show file tree

Hide file tree

Showing 6 changed files with 193 additions and 35 deletions.
diff --git a/crates/ruff_python_formatter/resources/test/fixtures/ruff/expression/set_comp.py b/crates/ruff_python_formatter/resources/test/fixtures/ruff/expression/set_comp.py
@@ -0,0 +1,45 @@
+{i for i in []}
+
+{i for i in [1,]}
+
+{
+    a  # a
+    for # for
+    c  # c
+    in  # in
+    e  # e
+}
+
+{
+    # above a
+    a  # a
+    # above for
+    for # for
+    # above c
+    c  # c
+    # above in
+    in  # in
+    # above e
+    e  # e
+    # above if
+    if  # if
+    # above f
+    f  # f
+    # above if2
+    if  # if2
+    # above g
+    g  # g
+}
+
+{
+    aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa + bbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbb + [dddddddddddddddddd, eeeeeeeeeeeeeeeeeee]
+    for
+    ccccccccccccccccccccccccccccccccccccccc,
+    ddddddddddddddddddd, [eeeeeeeeeeeeeeeeeeeeee, fffffffffffffffffffffffff]
+    in
+    eeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeffffffffffffffffffffffffggggggggggggggggggggghhhhhhhhhhhhhhothermoreeand_even_moreddddddddddddddddddddd
+    if
+    fffffffffffffffffffffffffffffffffffffffffff < gggggggggggggggggggggggggggggggggggggggggggggg < hhhhhhhhhhhhhhhhhhhhhhhhhh
+    if
+    gggggggggggggggggggggggggggggggggggggggggggg
+}
diff --git a/crates/ruff_python_formatter/src/expression/expr_set_comp.rs b/crates/ruff_python_formatter/src/expression/expr_set_comp.rs
@@ -1,19 +1,39 @@
 use crate::context::PyFormatContext;
-use crate::expression::parentheses::{NeedsParentheses, OptionalParentheses};
-use crate::{not_yet_implemented_custom_text, FormatNodeRule, PyFormatter};
-use ruff_formatter::{write, Buffer, FormatResult};
+use crate::expression::parentheses::{parenthesized, NeedsParentheses, OptionalParentheses};
+use crate::prelude::*;
+use crate::AsFormat;
+use crate::{FormatNodeRule, PyFormatter};
+use ruff_formatter::{format_args, write, Buffer, FormatResult};
 use ruff_python_ast::node::AnyNodeRef;
 use rustpython_parser::ast::ExprSetComp;
 
 #[derive(Default)]
 pub struct FormatExprSetComp;
 
 impl FormatNodeRule<ExprSetComp> for FormatExprSetComp {
-    fn fmt_fields(&self, _item: &ExprSetComp, f: &mut PyFormatter) -> FormatResult<()> {
+    fn fmt_fields(&self, item: &ExprSetComp, f: &mut PyFormatter) -> FormatResult<()> {
+        let ExprSetComp {
+            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_IMPLEMENTED_set_value for value in NOT_IMPLEMENTED_set}"
+            [parenthesized(
+                "{",
+                &format_args!(
+                    group(&elt.format()),
+                    soft_line_break_or_space(),
+                    group(&joined)
+                ),
+                "}"
             )]
         )
     }

diff --git a/crates/ruff_python_formatter/tests/snapshots/black_compatibility@py_39__pep_572_py39.py.snap b/crates/ruff_python_formatter/tests/snapshots/black_compatibility@py_39__pep_572_py39.py.snap
@@ -19,12 +19,8 @@ x[(a := 1), (b := 3)]
 ```diff
 --- Black
 +++ Ruff
-@@ -1,7 +1,7 @@
- # Unparenthesized walruses are now allowed in set literals & set comprehensions
- # since Python 3.9
- {x := 1, 2, 3}
--{x4 := x**5 for x in range(7)}
-+{NOT_IMPLEMENTED_set_value for value in NOT_IMPLEMENTED_set}
+@@ -4,4 +4,4 @@
+ {x4 := x**5 for x in range(7)}
  # We better not remove the parentheses here (since it's a 3.10 feature)
  x[(a := 1)]
 -x[(a := 1), (b := 3)]
@@ -37,7 +33,7 @@ x[(a := 1), (b := 3)]
 # Unparenthesized walruses are now allowed in set literals & set comprehensions
 # since Python 3.9
 {x := 1, 2, 3}
-{NOT_IMPLEMENTED_set_value for value in NOT_IMPLEMENTED_set}
+{x4 := x**5 for x in range(7)}
 # We better not remove the parentheses here (since it's a 3.10 feature)
 x[(a := 1)]
 x[((a := 1), (b := 3))]

diff --git a/...uff_python_formatter/tests/snapshots/black_compatibility@simple_cases__expression.py.snap b/...uff_python_formatter/tests/snapshots/black_compatibility@simple_cases__expression.py.snap
@@ -316,19 +316,7 @@ last_call()
  ()
  (1,)
  (1, 2)
-@@ -87,22 +90,19 @@
-     another,
-     *more,
- ]
--{i for i in (1, 2, 3)}
--{(i**2) for i in (1, 2, 3)}
--{(i**2) for i, _ in ((1, "a"), (2, "b"), (3, "c"))}
--{((i**2) + j) for i in (1, 2, 3) for j in (1, 2, 3)}
-+{NOT_IMPLEMENTED_set_value for value in NOT_IMPLEMENTED_set}
-+{NOT_IMPLEMENTED_set_value for value in NOT_IMPLEMENTED_set}
-+{NOT_IMPLEMENTED_set_value for value in NOT_IMPLEMENTED_set}
-+{NOT_IMPLEMENTED_set_value for value in NOT_IMPLEMENTED_set}
- [i for i in (1, 2, 3)]
+@@ -95,14 +98,11 @@
  [(i**2) for i in (1, 2, 3)]
  [(i**2) for i, _ in ((1, "a"), (2, "b"), (3, "c"))]
  [((i**2) + j) for i in (1, 2, 3) for j in (1, 2, 3)]
@@ -599,10 +587,10 @@ str or None if (1 if True else 2) else str or bytes or None
     another,
     *more,
 ]
-{NOT_IMPLEMENTED_set_value for value in NOT_IMPLEMENTED_set}
-{NOT_IMPLEMENTED_set_value for value in NOT_IMPLEMENTED_set}
-{NOT_IMPLEMENTED_set_value for value in NOT_IMPLEMENTED_set}
-{NOT_IMPLEMENTED_set_value for value in NOT_IMPLEMENTED_set}
+{i for i in (1, 2, 3)}
+{(i**2) for i in (1, 2, 3)}
+{(i**2) for i, _ in ((1, "a"), (2, "b"), (3, "c"))}
+{((i**2) + j) for i in (1, 2, 3) for j in (1, 2, 3)}
 [i for i in (1, 2, 3)]
 [(i**2) for i in (1, 2, 3)]
 [(i**2) for i, _ in ((1, "a"), (2, "b"), (3, "c"))]

diff --git a/crates/ruff_python_formatter/tests/snapshots/format@expression__binary.py.snap b/crates/ruff_python_formatter/tests/snapshots/format@expression__binary.py.snap
@@ -285,7 +285,10 @@ aaaaaaaaaaaaaa + [
     aaaaaaaaaaaaaa
     + (NOT_YET_IMPLEMENTED_generator_key for NOT_YET_IMPLEMENTED_generator_key in [])
 )
-aaaaaaaaaaaaaa + {NOT_IMPLEMENTED_set_value for value in NOT_IMPLEMENTED_set}
+aaaaaaaaaaaaaa + {
+    a
+    for x in bbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbb
+}
 
 # Wraps it in parentheses if it needs to break both left and right
 (
@@ -295,7 +298,10 @@ aaaaaaaaaaaaaa + {NOT_IMPLEMENTED_set_value for value in NOT_IMPLEMENTED_set}
 
 
 # But only for expressions that have a statement parent.
-not (aaaaaaaaaaaaaa + {NOT_IMPLEMENTED_set_value for value in NOT_IMPLEMENTED_set})
+not (
+    aaaaaaaaaaaaaa
+    + {a for x in bbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbb}
+)
 [
     a
     + [
@@ -476,9 +482,7 @@ if (
 
 
 # Unstable formatting in https://github.com/realtyem/synapse-unraid/blob/unraid_develop/synapse/handlers/presence.py
-for user_id in set(
-    target_user_ids
-) - {NOT_IMPLEMENTED_set_value for value in NOT_IMPLEMENTED_set}:
+for user_id in set(target_user_ids) - {u.user_id for u in updates}:
     updates.append(UserPresenceState.default(user_id))
 
 # Keeps parenthesized left hand sides

diff --git a/crates/ruff_python_formatter/tests/snapshots/format@expression__set_comp.py.snap b/crates/ruff_python_formatter/tests/snapshots/format@expression__set_comp.py.snap
@@ -0,0 +1,105 @@
+---
+source: crates/ruff_python_formatter/tests/fixtures.rs
+input_file: crates/ruff_python_formatter/resources/test/fixtures/ruff/expression/set_comp.py
+---
+## Input
+```py
+{i for i in []}
+
+{i for i in [1,]}
+
+{
+    a  # a
+    for # for
+    c  # c
+    in  # in
+    e  # e
+}
+
+{
+    # above a
+    a  # a
+    # above for
+    for # for
+    # above c
+    c  # c
+    # above in
+    in  # in
+    # above e
+    e  # e
+    # above if
+    if  # if
+    # above f
+    f  # f
+    # above if2
+    if  # if2
+    # above g
+    g  # g
+}
+
+{
+    aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa + bbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbb + [dddddddddddddddddd, eeeeeeeeeeeeeeeeeee]
+    for
+    ccccccccccccccccccccccccccccccccccccccc,
+    ddddddddddddddddddd, [eeeeeeeeeeeeeeeeeeeeee, fffffffffffffffffffffffff]
+    in
+    eeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeffffffffffffffffffffffffggggggggggggggggggggghhhhhhhhhhhhhhothermoreeand_even_moreddddddddddddddddddddd
+    if
+    fffffffffffffffffffffffffffffffffffffffffff < gggggggggggggggggggggggggggggggggggggggggggggg < hhhhhhhhhhhhhhhhhhhhhhhhhh
+    if
+    gggggggggggggggggggggggggggggggggggggggggggg
+}
+```
+
+## Output
+```py
+{i for i in []}
+
+{
+    i
+    for i in [
+        1,
+    ]
+}
+
+{
+    a  # a
+    for c in e  # for  # c  # in  # e
+}
+
+{
+    # above a
+    a  # a
+    # above for
+    for  # for
+    # above c
+    c  # c
+    # above in
+    in  # in
+    # above e
+    e  # e
+    # above if
+    if  # if
+    # above f
+    f  # f
+    # above if2
+    if  # if2
+    # above g
+    g  # g
+}
+
+{
+    aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa
+    + bbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbb
+    + [dddddddddddddddddd, eeeeeeeeeeeeeeeeeee]
+    for (
+        ccccccccccccccccccccccccccccccccccccccc,
+        ddddddddddddddddddd,
+        [eeeeeeeeeeeeeeeeeeeeee, fffffffffffffffffffffffff],
+    ) in eeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeffffffffffffffffffffffffggggggggggggggggggggghhhhhhhhhhhhhhothermoreeand_even_moreddddddddddddddddddddd
+    if fffffffffffffffffffffffffffffffffffffffffff
+    < gggggggggggggggggggggggggggggggggggggggggggggg
+    < hhhhhhhhhhhhhhhhhhhhhhhhhh
+    if gggggggggggggggggggggggggggggggggggggggggggg
+}
+```