Add support for special comments in multiline functions (#642)

black.py

@@ -2112,8 +2112,19 @@ def split_line(
         return
 
     line_str = str(line).strip("\n")
-    if not line.should_explode and is_line_short_enough(
-        line, line_length=line_length, line_str=line_str
+
+    # 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
+        and not line.should_explode
+        and is_line_short_enough(line, line_length=line_length, line_str=line_str)
     ):
         yield line
         return
@@ -2462,6 +2473,16 @@ def is_import(leaf: Leaf) -> bool:
     )
 
 
+def is_special_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:"))
+    )
+
+
 def normalize_prefix(leaf: Leaf, *, inside_brackets: bool) -> None:
     """Leave existing extra newlines if not `inside_brackets`. Remove everything
     else.
@@ -2951,6 +2972,7 @@ def ensure_visible(leaf: Leaf) -> None:
 
 def should_explode(line: Line, opening_bracket: Leaf) -> bool:
     """Should `line` immediately be split with `delimiter_split()` after RHS?"""
+
     if not (
         opening_bracket.parent
         and opening_bracket.parent.type in {syms.atom, syms.import_from}

tests/data/comments6.py

@@ -0,0 +1,68 @@
+from typing import Any, Tuple
+
+
+def f(
+    a,  # type: int
+):
+    pass
+
+
+# test type comments
+def f(a, b, c, d, e, f, g, h, i):
+    # type: (int, int, int, int, int, int, int, int, int) -> None
+    pass
+
+
+def f(
+    a,  # type: int
+    b,  # type: int
+    c,  # type: int
+    d,  # type: int
+    e,  # type: int
+    f,  # type: int
+    g,  # type: int
+    h,  # type: int
+    i,  # type: int
+):
+    # type: (...) -> None
+    pass
+
+
+def f(
+    arg,  # type: int
+    *args,  # type: *Any
+    default=False,  # type: bool
+    **kwargs,  # type: **Any
+):
+    # type: (...) -> None
+    pass
+
+
+def f(
+    a,  # type: int
+    b,  # type: int
+    c,  # type: int
+    d,  # type: int
+):
+    # type: (...) -> None
+
+    element = 0  # type: int
+    another_element = 1  # type: float
+    another_element_with_long_name = 2  # type: int
+    another_really_really_long_element_with_a_unnecessarily_long_name_to_describe_what_it_does_enterprise_style = (
+        3
+    )  # type: int
+
+    tup = (
+        another_element,  # type: int
+        another_really_really_long_element_with_a_unnecessarily_long_name_to_describe_what_it_does_enterprise_style,  # type: int
+    )  # type: Tuple[int, int]
+
+    a = (
+        element
+        + another_element
+        + another_element_with_long_name
+        + element
+        + another_element
+        + another_element_with_long_name
+    )  # type: int

tests/test_black.py

@@ -362,6 +362,14 @@ def test_comments5(self) -> None:
         black.assert_equivalent(source, actual)
         black.assert_stable(source, actual, line_length=ll)
 
+    @patch("black.dump_to_file", dump_to_stderr)
+    def test_comments6(self) -> None:
+        source, expected = read_data("comments6")
+        actual = fs(source)
+        self.assertFormatEqual(expected, actual)
+        black.assert_equivalent(source, actual)
+        black.assert_stable(source, actual, line_length=ll)
+
     @patch("black.dump_to_file", dump_to_stderr)
     def test_cantfit(self) -> None:
         source, expected = read_data("cantfit")