tweak collection literals to explode with trailing comma

black.py

@@ -1240,6 +1240,61 @@ def is_stub_class(self) -> bool:
             Leaf(token.DOT, ".") for _ in range(3)
         ]
 
+    @property
+    def is_collection_with_optional_trailing_comma(self) -> bool:
+        """Is this line a collection literal with a trailing comma that's optional?
+
+        Note that the trailing comma in a 1-tuple is not optional.
+        """
+        if not self.leaves or len(self.leaves) < 4:
+            return False
+        # Look for and address a trailing colon.
+        if self.leaves[-1].type == token.COLON:
+            closer = self.leaves[-2]
+            close_index = -2
+        else:
+            closer = self.leaves[-1]
+            close_index = -1
+        if closer.type not in CLOSING_BRACKETS or self.inside_brackets:
+            return False
+        if closer.type == token.RPAR:
+            # Tuples require an extra check, because if there's only
+            # one element in the tuple removing the comma unmakes the
+            # tuple.
+            #
+            # We also check for parens before looking for the trailing
+            # comma because in some cases (eg assigning a dict
+            # literal) the literal gets wrapped in temporary parens
+            # during parsing. This case is covered by the
+            # collections.py test data.
+            opener = closer.opening_bracket
+            for _open_index, leaf in enumerate(self.leaves):
+                if leaf is opener:
+                    break
+            else:
+                # Couldn't find the matching opening paren, play it safe.
+                return False
+            commas = 0
+            comma_depth = self.leaves[close_index - 1].bracket_depth
+            for leaf in self.leaves[_open_index + 1 : close_index]:
+                if leaf.bracket_depth == comma_depth and leaf.type == token.COMMA:
+                    commas += 1
+            if commas > 1:
+                # We haven't looked yet for the trailing comma because
+                # we might also have caught noop parens.
+                return self.leaves[close_index - 1].type == token.COMMA
+            elif commas == 1:
+                return False  # it's either a one-tuple or didn't have a trailing comma
+            if self.leaves[close_index - 1].type in CLOSING_BRACKETS:
+                close_index -= 1
+                closer = self.leaves[close_index]
+                if closer.type == token.RPAR:
+                    # TODO: this is a gut feeling. Will we ever see this?
+                    return False
+        if self.leaves[close_index - 1].type != token.COMMA:
+            return False
+        return True
+
     @property
     def is_def(self) -> bool:
         """Is this a function definition? (Also returns True for async defs.)"""
@@ -1356,60 +1411,39 @@ def contains_multiline_strings(self) -> bool:
 
     def maybe_remove_trailing_comma(self, closing: Leaf) -> bool:
         """Remove trailing comma if there is one and it's safe."""
+        if not (self.leaves and self.leaves[-1].type == token.COMMA):
+            return False
+        # We remove trailing commas only in the case of importing a
+        # single name from a module.
         if not (
             self.leaves
+            and self.is_import
+            and len(self.leaves) > 4
             and self.leaves[-1].type == token.COMMA
             and closing.type in CLOSING_BRACKETS
+            and self.leaves[-4].type == token.NAME
+            and (
+                # regular `from foo import bar,`
+                self.leaves[-4].value == "import"
+                # `from foo import (bar as baz,)
+                or (
+                    len(self.leaves) > 6
+                    and self.leaves[-6].value == "import"
+                    and self.leaves[-3].value == "as"
+                )
+                # `from foo import bar as baz,`
+                or (
+                    len(self.leaves) > 5
+                    and self.leaves[-5].value == "import"
+                    and self.leaves[-3].value == "as"
+                )
+            )
+            and closing.type == token.RPAR
         ):
             return False
 
-        if closing.type == token.RBRACE:
-            self.remove_trailing_comma()
-            return True
-
-        if closing.type == token.RSQB:
-            comma = self.leaves[-1]
-            if comma.parent and comma.parent.type == syms.listmaker:
-                self.remove_trailing_comma()
-                return True
-
-        # For parens let's check if it's safe to remove the comma.
-        # Imports are always safe.
-        if self.is_import:
-            self.remove_trailing_comma()
-            return True
-
-        # Otherwise, if the trailing one is the only one, we might mistakenly
-        # change a tuple into a different type by removing the comma.
-        depth = closing.bracket_depth + 1
-        commas = 0
-        opening = closing.opening_bracket
-        for _opening_index, leaf in enumerate(self.leaves):
-            if leaf is opening:
-                break
-
-        else:
-            return False
-
-        for leaf in self.leaves[_opening_index + 1 :]:
-            if leaf is closing:
-                break
-
-            bracket_depth = leaf.bracket_depth
-            if bracket_depth == depth and leaf.type == token.COMMA:
-                commas += 1
-                if leaf.parent and leaf.parent.type in {
-                    syms.arglist,
-                    syms.typedargslist,
-                }:
-                    commas += 1
-                    break
-
-        if commas > 1:
-            self.remove_trailing_comma()
-            return True
-
-        return False
+        self.remove_trailing_comma()
+        return True
 
     def append_comment(self, comment: Leaf) -> bool:
         """Add an inline or standalone comment to the line."""
@@ -2354,6 +2388,7 @@ def split_line(
     if (
         not line.contains_uncollapsable_type_comments()
         and not line.should_explode
+        and not line.is_collection_with_optional_trailing_comma
         and (
             is_line_short_enough(line, line_length=line_length, line_str=line_str)
             or line.contains_unsplittable_type_ignore()

tests/data/collections.py

@@ -0,0 +1,160 @@
+import core, time, a
+
+from . import A, B, C
+
+# unwraps
+from foo import (
+    bar,
+)
+
+# stays wrapped
+from foo import (
+    baz,
+    qux,
+)
+
+# as doesn't get confusing when unwrapped
+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
+
+# unwraps
+from foo import bar
+
+# stays wrapped
+from foo import (
+    baz,
+    qux,
+)
+
+# as doesn't get confusing when unwrapped
+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,},
+    )

tests/data/comments2.py

@@ -63,7 +63,7 @@ def inline_comments_in_brackets_ruin_everything():
         parameters.children = [
             children[0],  # (1
             body,
-            children[-1],  # )1
+            children[-1]  # )1
         ]
         parameters.children = [
             children[0],
@@ -142,7 +142,7 @@ def inline_comments_in_brackets_ruin_everything():
         syms.simple_stmt,
         [
             Node(statement, result),
-            Leaf(token.NEWLINE, '\n'),  # FIXME: \r\n?
+            Leaf(token.NEWLINE, '\n')  # FIXME: \r\n?
         ],
     )
 

tests/data/comments7.py

@@ -94,7 +94,7 @@ def func():
 
 def func():
     c = call(
-        0.0123, 0.0456, 0.0789, 0.0123, 0.0789, a[-1]  # type: ignore
+        0.0123, 0.0456, 0.0789, 0.0123, 0.0789, a[-1],  # type: ignore
     )
 
     # The type: ignore exception only applies to line length, not