Fix merging implicit multiline strings that have inline comments (#3956)

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

diff --git a/src/black/trans.py b/src/black/trans.py
index daed26427d79f2518cec3df0ff32f6e82025506e..a3f6467cc9e3920758ce3b5ddf00a5e2c41129b6 100644 (file)
--- a/src/black/trans.py
+++ b/src/black/trans.py
@@ -1,6 +1,7 @@
 """
 String transformers that can split and merge strings.
 """
+
 import re
 from abc import ABC, abstractmethod
 from collections import defaultdict
@@ -389,7 +390,19 @@ class StringMerger(StringTransformer, CustomSplitMapMixin):
                 and is_valid_index(idx + 1)
                 and LL[idx + 1].type == token.STRING
             ):
-                if not is_part_of_annotation(leaf):
+                # Let's check if the string group contains an inline comment
+                # If we have a comment inline, we don't merge the strings
+                contains_comment = False
+                i = idx
+                while is_valid_index(i):
+                    if LL[i].type != token.STRING:
+                        break
+                    if line.comments_after(LL[i]):
+                        contains_comment = True
+                        break
+                    i += 1
+
+                if not is_part_of_annotation(leaf) and not contains_comment:
                     string_indices.append(idx)
 
                 # Advance to the next non-STRING leaf.
@@ -942,6 +955,9 @@ class StringParenStripper(StringTransformer):
                 LL[lpar_or_rpar_idx].remove()  # Remove lpar.
                 replace_child(LL[idx], string_leaf)
                 new_line.append(string_leaf)
+                # replace comments
+                old_comments = new_line.comments.pop(id(LL[idx]), [])
+                new_line.comments.setdefault(id(string_leaf), []).extend(old_comments)
             else:
                 LL[lpar_or_rpar_idx].remove()  # This is a rpar.