Actually add trailing commas to collection literals even if there are terminating comments (#3393)

Co-authored-by: Jelle Zijlstra <jelle.zijlstra@gmail.com>
Co-authored-by: Richard Si <sichard26@gmail.com>

CHANGES.md

@@ -14,6 +14,9 @@
 
 <!-- Changes that affect Black's preview style -->
 
+- Add trailing commas to collection literals even if there's a comment after the last
+  entry (#3393)
+
 ### Configuration
 
 <!-- Changes to how Black can be configured -->

src/black/linegen.py

@@ -520,7 +520,7 @@ def transform_line(
     else:
 
         def _rhs(
-            self: object, line: Line, features: Collection[Feature]
+            self: object, line: Line, features: Collection[Feature], mode: Mode
         ) -> Iterator[Line]:
             """Wraps calls to `right_hand_split`.
 
@@ -604,7 +604,9 @@ class _BracketSplitComponent(Enum):
     tail = auto()
 
 
-def left_hand_split(line: Line, _features: Collection[Feature] = ()) -> Iterator[Line]:
+def left_hand_split(
+    line: Line, _features: Collection[Feature], mode: Mode
+) -> Iterator[Line]:
     """Split line into many lines, starting with the first matching bracket pair.
 
     Note: this usually looks weird, only use this for function definitions.
@@ -940,16 +942,39 @@ def dont_increase_indentation(split_func: Transformer) -> Transformer:
     """
 
     @wraps(split_func)
-    def split_wrapper(line: Line, features: Collection[Feature] = ()) -> Iterator[Line]:
+    def split_wrapper(
-        for split_line in split_func(line, features):
+        line: Line, features: Collection[Feature], mode: Mode
+    ) -> Iterator[Line]:
+        for split_line in split_func(line, features, mode):
             normalize_prefix(split_line.leaves[0], inside_brackets=True)
             yield split_line
 
     return split_wrapper
 
 
+def _get_last_non_comment_leaf(line: Line) -> Optional[int]:
+    for leaf_idx in range(len(line.leaves) - 1, 0, -1):
+        if line.leaves[leaf_idx].type != STANDALONE_COMMENT:
+            return leaf_idx
+    return None
+
+
+def _safe_add_trailing_comma(safe: bool, delimiter_priority: int, line: Line) -> Line:
+    if (
+        safe
+        and delimiter_priority == COMMA_PRIORITY
+        and line.leaves[-1].type != token.COMMA
+        and line.leaves[-1].type != STANDALONE_COMMENT
+    ):
+        new_comma = Leaf(token.COMMA, ",")
+        line.append(new_comma)
+    return line
+
+
 @dont_increase_indentation
-def delimiter_split(line: Line, features: Collection[Feature] = ()) -> Iterator[Line]:
+def delimiter_split(
+    line: Line, features: Collection[Feature], mode: Mode
+) -> Iterator[Line]:
     """Split according to delimiters of the highest priority.
 
     If the appropriate Features are given, the split will add trailing commas
@@ -989,7 +1014,8 @@ def append_to_line(leaf: Leaf) -> Iterator[Line]:
             )
             current_line.append(leaf)
 
-    for leaf in line.leaves:
+    last_non_comment_leaf = _get_last_non_comment_leaf(line)
+    for leaf_idx, leaf in enumerate(line.leaves):
         yield from append_to_line(leaf)
 
         for comment_after in line.comments_after(leaf):
@@ -1006,6 +1032,15 @@ def append_to_line(leaf: Leaf) -> Iterator[Line]:
                     trailing_comma_safe and Feature.TRAILING_COMMA_IN_CALL in features
                 )
 
+        if (
+            Preview.add_trailing_comma_consistently in mode
+            and last_leaf.type == STANDALONE_COMMENT
+            and leaf_idx == last_non_comment_leaf
+        ):
+            current_line = _safe_add_trailing_comma(
+                trailing_comma_safe, delimiter_priority, current_line
+            )
+
         leaf_priority = bt.delimiters.get(id(leaf))
         if leaf_priority == delimiter_priority:
             yield current_line
@@ -1014,20 +1049,15 @@ def append_to_line(leaf: Leaf) -> Iterator[Line]:
                 mode=line.mode, depth=line.depth, inside_brackets=line.inside_brackets
             )
     if current_line:
-        if (
+        current_line = _safe_add_trailing_comma(
-            trailing_comma_safe
+            trailing_comma_safe, delimiter_priority, current_line
-            and delimiter_priority == COMMA_PRIORITY
+        )
-            and current_line.leaves[-1].type != token.COMMA
-            and current_line.leaves[-1].type != STANDALONE_COMMENT
-        ):
-            new_comma = Leaf(token.COMMA, ",")
-            current_line.append(new_comma)
         yield current_line
 
 
 @dont_increase_indentation
 def standalone_comment_split(
-    line: Line, features: Collection[Feature] = ()
+    line: Line, features: Collection[Feature], mode: Mode
 ) -> Iterator[Line]:
     """Split standalone comments from the rest of the line."""
     if not line.contains_standalone_comments(0):
@@ -1480,7 +1510,7 @@ def run_transformer(
     if not line_str:
         line_str = line_to_string(line)
     result: List[Line] = []
-    for transformed_line in transform(line, features):
+    for transformed_line in transform(line, features, mode):
         if str(transformed_line).strip("\n") == line_str:
             raise CannotTransform("Line transformer returned an unchanged result")
 

src/black/mode.py

@@ -153,6 +153,7 @@ def supports_feature(target_versions: Set[TargetVersion], feature: Feature) -> b
 class Preview(Enum):
     """Individual preview style features."""
 
+    add_trailing_comma_consistently = auto()
     hex_codes_in_unicode_sequences = auto()
     prefer_splitting_right_hand_side_of_assignments = auto()
     # NOTE: string_processing requires wrap_long_dict_values_in_parens

src/black/trans.py

@@ -32,7 +32,7 @@
 
 from black.comments import contains_pragma_comment
 from black.lines import Line, append_leaves
-from black.mode import Feature
+from black.mode import Feature, Mode
 from black.nodes import (
     CLOSING_BRACKETS,
     OPENING_BRACKETS,
@@ -63,7 +63,7 @@ class CannotTransform(Exception):
 # types
 T = TypeVar("T")
 LN = Union[Leaf, Node]
-Transformer = Callable[[Line, Collection[Feature]], Iterator[Line]]
+Transformer = Callable[[Line, Collection[Feature], Mode], Iterator[Line]]
 Index = int
 NodeType = int
 ParserState = int
@@ -81,7 +81,9 @@ def TErr(err_msg: str) -> Err[CannotTransform]:
     return Err(cant_transform)
 
 
-def hug_power_op(line: Line, features: Collection[Feature]) -> Iterator[Line]:
+def hug_power_op(
+    line: Line, features: Collection[Feature], mode: Mode
+) -> Iterator[Line]:
     """A transformer which normalizes spacing around power operators."""
 
     # Performance optimization to avoid unnecessary Leaf clones and other ops.
@@ -228,7 +230,9 @@ def do_transform(
             yield an CannotTransform after that point.)
         """
 
-    def __call__(self, line: Line, _features: Collection[Feature]) -> Iterator[Line]:
+    def __call__(
+        self, line: Line, _features: Collection[Feature], _mode: Mode
+    ) -> Iterator[Line]:
         """
         StringTransformer instances have a call signature that mirrors that of
         the Transformer type.

tests/data/preview/trailing_comma.py

@@ -0,0 +1,55 @@
+e = {
+    "a": fun(msg, "ts"),
+    "longggggggggggggggid": ...,
+    "longgggggggggggggggggggkey": ..., "created": ...
+    # "longkey": ...
+}
+f = [
+    arg1,
+    arg2,
+    arg3, arg4
+    # comment
+]
+g = (
+    arg1,
+    arg2,
+    arg3, arg4
+    # comment
+)
+h = {
+    arg1,
+    arg2,
+    arg3, arg4
+    # comment
+}
+
+# output
+
+e = {
+    "a": fun(msg, "ts"),
+    "longggggggggggggggid": ...,
+    "longgggggggggggggggggggkey": ...,
+    "created": ...,
+    # "longkey": ...
+}
+f = [
+    arg1,
+    arg2,
+    arg3,
+    arg4,
+    # comment
+]
+g = (
+    arg1,
+    arg2,
+    arg3,
+    arg4,
+    # comment
+)
+h = {
+    arg1,
+    arg2,
+    arg3,
+    arg4,
+    # comment
+}

tests/data/preview_context_managers/targeting_py39.py

@@ -84,7 +84,7 @@
     # First comment.
     new_new_new1() as cm1,
     # Second comment.
-    new_new_new2()
+    new_new_new2(),
     # Last comment.
 ):
     pass