Copy over comments when hugging power ops (#2874)

[etc/vim.git] / src / black / trans.py

diff --git a/src/black/trans.py b/src/black/trans.py
index 01aa80eaaf84760f765ae34eef651d7336e28c09..28d9250adc1fe78d7e788b2d6fd382d00d42073b 100644 (file)
--- a/src/black/trans.py
+++ b/src/black/trans.py
@@ -121,7 +121,7 @@ def hug_power_op(line: Line, features: Collection[Feature]) -> Iterator[Line]:
 
         return False
 
 
         return False
 

-    leaves: List[Leaf] = []
+    new_line = line.clone()

     should_hug = False
     for idx, leaf in enumerate(line.leaves):
         new_leaf = leaf.clone()
     should_hug = False
     for idx, leaf in enumerate(line.leaves):
         new_leaf = leaf.clone()


@@ -139,18 +139,14 @@ def hug_power_op(line: Line, features: Collection[Feature]) -> Iterator[Line]:
         if should_hug:
             new_leaf.prefix = ""
 
         if should_hug:
             new_leaf.prefix = ""
 

-        leaves.append(new_leaf)
-
-    yield Line(
-        mode=line.mode,
-        depth=line.depth,
-        leaves=leaves,
-        comments=line.comments,
-        bracket_tracker=line.bracket_tracker,
-        inside_brackets=line.inside_brackets,
-        should_split_rhs=line.should_split_rhs,
-        magic_trailing_comma=line.magic_trailing_comma,
-    )
+        # We have to be careful to make a new line properly:
+        # - bracket related metadata must be maintained (handled by Line.append)
+        # - comments need to copied over, updating the leaf IDs they're attached to
+        new_line.append(new_leaf, preformatted=True)
+        for comment_leaf in line.comments_after(leaf):
+            new_line.append(comment_leaf, preformatted=True)
+
+    yield new_line

 
 
 class StringTransformer(ABC):
 
 
 class StringTransformer(ABC):