[2097] Fix is_type_comment

src/black/nodes.py

@@ -2,6 +2,7 @@
 blib2to3 Node/Leaf transformation-related utility functions.
 """
 
+import re
 import sys
 from typing import Generic, Iterator, List, Optional, Set, Tuple, TypeVar, Union
 
@@ -778,12 +779,21 @@ def is_import(leaf: Leaf) -> bool:
     )
 
 
+is_type_comment_regex = re.compile(r'^#[\s\t]+type[\s\t]*:')
+
+
 def is_type_comment(leaf: Leaf, suffix: str = "") -> 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 t in {token.COMMENT, STANDALONE_COMMENT} and v.startswith("# type:" + suffix)
+    if t not in {token.COMMENT, STANDALONE_COMMENT}:
+        return False
+
+    if not v.startswith("#"):
+        return False
+
+    return bool(is_type_comment_regex.match(v))
 
 
 def wrap_in_parentheses(parent: Node, child: LN, *, visible: bool = True) -> None:

tests/test_typed_comments.py

@@ -0,0 +1,13 @@
+import pytest
+
+from black.nodes import is_type_comment
+from blib2to3.pgen2 import token
+from blib2to3.pytree import Leaf
+
+
+@pytest.mark.parametrize(
+    "comment_str",
+    ["# type: int", "#  type: int", "#   type: int", "# type : int", "# type  : int"],
+)
+def test_typed_comments(comment_str: str) -> None:
+    assert is_type_comment(Leaf(token.COMMENT, comment_str))