Fix comment handling when parenthesising conditional expressions (#4134)

Fixes #3555

CHANGES.md

@@ -14,6 +14,7 @@
 
 <!-- Changes that affect Black's preview style -->
 
+- Fix comment handling when parenthesising conditional expressions (#4134)
 - Format module docstrings the same as class and function docstrings (#4095)
 - Fix bug where spaces were not added around parenthesized walruses in subscripts,
   unlike other binary operators (#4109)

src/black/linegen.py

@@ -170,8 +170,12 @@ def visit_test(self, node: Node) -> Iterator[Line]:
             )
 
             if not already_parenthesized:
+                # Similar to logic in wrap_in_parentheses
                 lpar = Leaf(token.LPAR, "")
                 rpar = Leaf(token.RPAR, "")
+                prefix = node.prefix
+                node.prefix = ""
+                lpar.prefix = prefix
                 node.insert_child(0, lpar)
                 node.append_child(rpar)
 

tests/data/cases/conditional_expression.py

@@ -67,6 +67,28 @@ def something():
         else ValuesListIterable
     )
 
+
+def foo(wait: bool = True):
+    # This comment is two
+    # lines long
+
+    # This is only one
+    time.sleep(1) if wait else None
+    time.sleep(1) if wait else None
+
+    # With newline above
+    time.sleep(1) if wait else None
+    # Without newline above
+    time.sleep(1) if wait else None
+
+
+a = "".join(
+    (
+        "",  # comment
+        "" if True else "",
+    )
+)
+
 # output
 
 long_kwargs_single_line = my_function(
@@ -159,3 +181,23 @@ def something():
         if named
         else FlatValuesListIterable if flat else ValuesListIterable
     )
+
+
+def foo(wait: bool = True):
+    # This comment is two
+    # lines long
+
+    # This is only one
+    time.sleep(1) if wait else None
+    time.sleep(1) if wait else None
+
+    # With newline above
+    time.sleep(1) if wait else None
+    # Without newline above
+    time.sleep(1) if wait else None
+
+
+a = "".join((
+    "",  # comment
+    "" if True else "",
+))