Enhance the type comment patch

README.md

@@ -948,6 +948,8 @@ More details can be found in [CONTRIBUTING](CONTRIBUTING.md).
 
 * long `del` statements are now split into multiple lines (#698)
 
+* type comments are no longer mangled in function signatures
+
 * improved performance of formatting deeply nested data structures (#509)
 
 * `blackd` now supports CORS (#622)
@@ -1370,5 +1372,6 @@ Multiple contributions by:
 * [Peter Bengtsson](mailto:mail@peterbe.com)
 * [Stavros Korokithakis](mailto:hi@stavros.io)
 * [Sunil Kapil](mailto:snlkapil@gmail.com)
+* [Utsav Shah](mailto:ukshah2@illinois.edu)
 * [Vishwas B Sharma](mailto:sharma.vishwas88@gmail.com)
 * [Chuck Wooters](mailto:chuck.wooters@microsoft.com)

black.py

@@ -1192,6 +1192,29 @@ def contains_standalone_comments(self, depth_limit: int = sys.maxsize) -> bool:
             if leaf.type == STANDALONE_COMMENT:
                 if leaf.bracket_depth <= depth_limit:
                     return True
+        return False
+
+    def contains_inner_type_comments(self) -> bool:
+        ignored_ids = set()
+        try:
+            last_leaf = self.leaves[-1]
+            ignored_ids.add(id(last_leaf))
+            if last_leaf.type == token.COMMA:
+                # When trailing commas are inserted by Black for consistency, comments
+                # after the previous last element are not moved (they don't have to,
+                # rendering will still be correct).  So we ignore trailing commas.
+                last_leaf = self.leaves[-2]
+                ignored_ids.add(id(last_leaf))
+        except IndexError:
+            return False
+
+        for leaf_id, comments in self.comments.items():
+            if leaf_id in ignored_ids:
+                continue
+
+            for comment in comments:
+                if is_type_comment(comment):
+                    return True
 
         return False
 
@@ -2135,16 +2158,8 @@ def split_line(
 
     line_str = str(line).strip("\n")
 
-    # we don't want to split special comments like type annotations
-    # https://github.com/python/typing/issues/186
-    has_special_comment = False
-    for leaf in line.leaves:
-        for comment in line.comments_after(leaf):
-            if leaf.type == token.COMMA and is_special_comment(comment):
-                has_special_comment = True
-
     if (
-        not has_special_comment
+        not line.contains_inner_type_comments()
         and not line.should_explode
         and is_line_short_enough(line, line_length=line_length, line_str=line_str)
     ):
@@ -2520,14 +2535,12 @@ def is_import(leaf: Leaf) -> bool:
     )
 
 
-def is_special_comment(leaf: Leaf) -> bool:
+def is_type_comment(leaf: Leaf) -> bool:
     """Return True if the given leaf is a special comment.
     Only returns true for type comments for now."""
     t = leaf.type
     v = leaf.value
-    return bool(
-        (t == token.COMMENT or t == STANDALONE_COMMENT) and (v.startswith("# type:"))
-    )
+    return t in {token.COMMENT, t == STANDALONE_COMMENT} and v.startswith("# type:")
 
 
 def normalize_prefix(leaf: Leaf, *, inside_brackets: bool) -> None:

tests/data/comments2.py

@@ -65,12 +65,18 @@ def inline_comments_in_brackets_ruin_everything():
             body,
             children[-1],  # )1
         ]
+        parameters.children = [
+            children[0],
+            body,
+            children[-1],  # type: ignore
+        ]
     else:
         parameters.children = [
             parameters.children[0],  # (2 what if this was actually long
             body,
             parameters.children[-1],  # )2
         ]
+        parameters.children = [parameters.what_if_this_was_actually_long.children[0], body, parameters.children[-1]]  # type: ignore
     if (self._proc is not None
             # has the child process finished?
             and self._returncode is None
@@ -217,12 +223,22 @@ def inline_comments_in_brackets_ruin_everything():
 def inline_comments_in_brackets_ruin_everything():
     if typedargslist:
         parameters.children = [children[0], body, children[-1]]  # (1  # )1
+        parameters.children = [
+            children[0],
+            body,
+            children[-1],  # type: ignore
+        ]
     else:
         parameters.children = [
             parameters.children[0],  # (2 what if this was actually long
             body,
             parameters.children[-1],  # )2
         ]
+        parameters.children = [
+            parameters.what_if_this_was_actually_long.children[0],
+            body,
+            parameters.children[-1],
+        ]  # type: ignore
     if (
         self._proc is not None
         # has the child process finished?

tests/data/comments6.py

@@ -52,6 +52,7 @@ def f(
     another_really_really_long_element_with_a_unnecessarily_long_name_to_describe_what_it_does_enterprise_style = (
         3
     )  # type: int
+    an_element_with_a_long_value = calls() or more_calls() and more()  # type: bool
 
     tup = (
         another_element,  # type: int
@@ -66,3 +67,20 @@ def f(
         + another_element
         + another_element_with_long_name
     )  # type: int
+
+
+def func(
+    a=some_list[0],  # type: int
+):  # type: () -> int
+    c = call(
+        0.0123,
+        0.0456,
+        0.0789,
+        0.0123,
+        0.0456,
+        0.0789,
+        0.0123,
+        0.0456,
+        0.0789,
+        a[-1],  # type: ignore
+    )