Add trailing commas to support collections.py test (#7)

astral-sh · Feb 15, 2023 · 8081cf5 · 8081cf5
1 parent 41c10f8
commit 8081cf5
Show file tree

Hide file tree

Showing 6 changed files with 323 additions and 123 deletions.
diff --git a/crates/ruff_fmt/resources/test/fixtures/black/simple_cases/collections.bak b/crates/ruff_fmt/resources/test/fixtures/black/simple_cases/collections.bak
@@ -0,0 +1,174 @@
+import core, time, a
+
+from . import A, B, C
+
+# keeps existing trailing comma
+from foo import (
+    bar,
+)
+
+# also keeps existing structure
+from foo import (
+    baz,
+    qux,
+)
+
+# `as` works as well
+from foo import (
+    xyzzy as magic,
+)
+
+a = {1,2,3,}
+b = {
+1,2,
+     3}
+c = {
+    1,
+    2,
+    3,
+}
+x = 1,
+y = narf(),
+nested = {(1,2,3),(4,5,6),}
+nested_no_trailing_comma = {(1,2,3),(4,5,6)}
+nested_long_lines = ["aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa", "bbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbb", "cccccccccccccccccccccccccccccccccccccccc", (1, 2, 3), "dddddddddddddddddddddddddddddddddddddddd"]
+{"oneple": (1,),}
+{"oneple": (1,)}
+['ls', 'lsoneple/%s' % (foo,)]
+x = {"oneple": (1,)}
+y = {"oneple": (1,),}
+assert False, ("aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa wraps %s" % bar)
+
+# looping over a 1-tuple should also not get wrapped
+for x in (1,):
+    pass
+for (x,) in (1,), (2,), (3,):
+    pass
+
+[1, 2, 3,]
+
+division_result_tuple = (6/2,)
+print("foo %r", (foo.bar,))
+
+if True:
+    IGNORED_TYPES_FOR_ATTRIBUTE_CHECKING = (
+        Config.IGNORED_TYPES_FOR_ATTRIBUTE_CHECKING
+        | {pylons.controllers.WSGIController}
+    )
+
+if True:
+    ec2client.get_waiter('instance_stopped').wait(
+        InstanceIds=[instance.id],
+        WaiterConfig={
+            'Delay': 5,
+        })
+    ec2client.get_waiter("instance_stopped").wait(
+        InstanceIds=[instance.id],
+        WaiterConfig={"Delay": 5,},
+    )
+    ec2client.get_waiter("instance_stopped").wait(
+        InstanceIds=[instance.id], WaiterConfig={"Delay": 5,},
+    )
+
+# output
+
+
+import core, time, a
+
+from . import A, B, C
+
+# keeps existing trailing comma
+from foo import (
+    bar,
+)
+
+# also keeps existing structure
+from foo import (
+    baz,
+    qux,
+)
+
+# `as` works as well
+from foo import (
+    xyzzy as magic,
+)
+
+a = {
+    1,
+    2,
+    3,
+}
+b = {1, 2, 3}
+c = {
+    1,
+    2,
+    3,
+}
+x = (1,)
+y = (narf(),)
+nested = {
+    (1, 2, 3),
+    (4, 5, 6),
+}
+nested_no_trailing_comma = {(1, 2, 3), (4, 5, 6)}
+nested_long_lines = [
+    "aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa",
+    "bbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbb",
+    "cccccccccccccccccccccccccccccccccccccccc",
+    (1, 2, 3),
+    "dddddddddddddddddddddddddddddddddddddddd",
+]
+{
+    "oneple": (1,),
+}
+{"oneple": (1,)}
+["ls", "lsoneple/%s" % (foo,)]
+x = {"oneple": (1,)}
+y = {
+    "oneple": (1,),
+}
+assert False, (
+    "aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa wraps %s"
+    % bar
+)
+
+# looping over a 1-tuple should also not get wrapped
+for x in (1,):
+    pass
+for (x,) in (1,), (2,), (3,):
+    pass
+
+[
+    1,
+    2,
+    3,
+]
+
+division_result_tuple = (6 / 2,)
+print("foo %r", (foo.bar,))
+
+if True:
+    IGNORED_TYPES_FOR_ATTRIBUTE_CHECKING = (
+        Config.IGNORED_TYPES_FOR_ATTRIBUTE_CHECKING
+        | {pylons.controllers.WSGIController}
+    )
+
+if True:
+    ec2client.get_waiter("instance_stopped").wait(
+        InstanceIds=[instance.id],
+        WaiterConfig={
+            "Delay": 5,
+        },
+    )
+    ec2client.get_waiter("instance_stopped").wait(
+        InstanceIds=[instance.id],
+        WaiterConfig={
+            "Delay": 5,
+        },
+    )
+    ec2client.get_waiter("instance_stopped").wait(
+        InstanceIds=[instance.id],
+        WaiterConfig={
+            "Delay": 5,
+        },
+    )
diff --git a/crates/ruff_fmt/resources/test/fixtures/black/simple_cases/collections.py b/crates/ruff_fmt/resources/test/fixtures/black/simple_cases/collections.py
@@ -69,106 +69,3 @@
     ec2client.get_waiter("instance_stopped").wait(
         InstanceIds=[instance.id], WaiterConfig={"Delay": 5,},
     )
-
-# output
-
-
-import core, time, a
-
-from . import A, B, C
-
-# keeps existing trailing comma
-from foo import (
-    bar,
-)
-
-# also keeps existing structure
-from foo import (
-    baz,
-    qux,
-)
-
-# `as` works as well
-from foo import (
-    xyzzy as magic,
-)
-
-a = {
-    1,
-    2,
-    3,
-}
-b = {1, 2, 3}
-c = {
-    1,
-    2,
-    3,
-}
-x = (1,)
-y = (narf(),)
-nested = {
-    (1, 2, 3),
-    (4, 5, 6),
-}
-nested_no_trailing_comma = {(1, 2, 3), (4, 5, 6)}
-nested_long_lines = [
-    "aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa",
-    "bbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbb",
-    "cccccccccccccccccccccccccccccccccccccccc",
-    (1, 2, 3),
-    "dddddddddddddddddddddddddddddddddddddddd",
-]
-{
-    "oneple": (1,),
-}
-{"oneple": (1,)}
-["ls", "lsoneple/%s" % (foo,)]
-x = {"oneple": (1,)}
-y = {
-    "oneple": (1,),
-}
-assert False, (
-    "aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa wraps %s"
-    % bar
-)
-
-# looping over a 1-tuple should also not get wrapped
-for x in (1,):
-    pass
-for (x,) in (1,), (2,), (3,):
-    pass
-
-[
-    1,
-    2,
-    3,
-]
-
-division_result_tuple = (6 / 2,)
-print("foo %r", (foo.bar,))
-
-if True:
-    IGNORED_TYPES_FOR_ATTRIBUTE_CHECKING = (
-        Config.IGNORED_TYPES_FOR_ATTRIBUTE_CHECKING
-        | {pylons.controllers.WSGIController}
-    )
-
-if True:
-    ec2client.get_waiter("instance_stopped").wait(
-        InstanceIds=[instance.id],
-        WaiterConfig={
-            "Delay": 5,
-        },
-    )
-    ec2client.get_waiter("instance_stopped").wait(
-        InstanceIds=[instance.id],
-        WaiterConfig={
-            "Delay": 5,
-        },
-    )
-    ec2client.get_waiter("instance_stopped").wait(
-        InstanceIds=[instance.id],
-        WaiterConfig={
-            "Delay": 5,
-        },
-    )
diff --git a/crates/ruff_fmt/src/format/expr.rs b/crates/ruff_fmt/src/format/expr.rs
@@ -229,7 +229,6 @@ fn format_list(
 ) -> FormatResult<()> {
     write!(f, [text("[")])?;
     if !elts.is_empty() {
-        // TODO(charlie): DRY.
         let magic_trailing_comma = expr
             .trivia
             .iter()
@@ -268,10 +267,16 @@ fn format_set(
     } else {
         write!(f, [text("{")])?;
         if !elts.is_empty() {
-            // TODO(charlie): DRY.
+            let magic_trailing_comma = expr
+                .trivia
+                .iter()
+                .any(|c| matches!(c.kind, TriviaKind::MagicTrailingComma));
             write!(
                 f,
                 [group(&format_args![soft_block_indent(&format_with(|f| {
+                    if magic_trailing_comma {
+                        write!(f, [expand_parent()])?;
+                    }
                     for (i, elt) in elts.iter().enumerate() {
                         write!(f, [group(&format_args![elt.format()])])?;
                         if i < elts.len() - 1 {
@@ -303,7 +308,7 @@ fn format_call(
         write!(f, [text(")")])?;
     } else {
         write!(f, [text("(")])?;
-        // TODO(charlie): DRY.
+
         let magic_trailing_comma = expr
             .trivia
             .iter()
@@ -576,10 +581,16 @@ fn format_dict(
 ) -> FormatResult<()> {
     write!(f, [text("{")])?;
     if !keys.is_empty() {
-        // TODO(charlie): DRY.
+        let magic_trailing_comma = expr
+            .trivia
+            .iter()
+            .any(|c| matches!(c.kind, TriviaKind::MagicTrailingComma));
         write!(
             f,
             [soft_block_indent(&format_with(|f| {
+                if magic_trailing_comma {
+                    write!(f, [expand_parent()])?;
+                }
                 for (i, (k, v)) in keys.iter().zip(values).enumerate() {
                     if let Some(k) = k {
                         write!(f, [k.format()])?;

diff --git a/crates/ruff_fmt/src/format/stmt.rs b/crates/ruff_fmt/src/format/stmt.rs
@@ -435,21 +435,28 @@ fn format_assert(
     msg: Option<&Expr>,
 ) -> FormatResult<()> {
     write!(f, [text("assert"), space()])?;
-    if is_self_closing(test) {
-        write!(f, [test.format()])?;
-    } else {
+    write!(
+        f,
+        [group(&format_args![
+            if_group_breaks(&text("(")),
+            soft_block_indent(&test.format()),
+            if_group_breaks(&text(")")),
+        ])]
+    )?;
+    if let Some(msg) = msg {
         write!(
             f,
-            [group(&format_args![
-                if_group_breaks(&text("(")),
-                soft_block_indent(&test.format()),
-                if_group_breaks(&text(")")),
-            ])]
+            [
+                text(","),
+                space(),
+                group(&format_args![
+                    if_group_breaks(&text("(")),
+                    soft_block_indent(&msg.format()),
+                    if_group_breaks(&text(")")),
+                ])
+            ]
         )?;
     }
-    if let Some(msg) = msg {
-        write!(f, [text(", "), msg.format()])?;
-    }
     Ok(())
 }
 
@@ -505,11 +512,18 @@ fn format_import_from(
     write!(f, [text("import")])?;
     write!(f, [space()])?;
 
+    let magic_trailing_comma = stmt
+        .trivia
+        .iter()
+        .any(|c| matches!(c.kind, TriviaKind::MagicTrailingComma));
     write!(
         f,
         [group(&format_args![
             if_group_breaks(&text("(")),
             soft_block_indent(&format_with(|f| {
+                if magic_trailing_comma {
+                    write!(f, [expand_parent()])?;
+                }
                 for (i, name) in names.iter().enumerate() {
                     write!(f, [name.format()])?;
                     if i < names.len() - 1 {