Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

[2097] Fix is_type_comment #3339

Open
wants to merge 12 commits into
base: main
Choose a base branch
from
9 changes: 8 additions & 1 deletion src/black/nodes.py
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand Down Expand Up @@ -778,12 +779,18 @@ 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} or not v.startswith("#"):
return False
itxasos23 marked this conversation as resolved.
Show resolved Hide resolved

return bool(is_type_comment_regex.match(v))


def wrap_in_parentheses(parent: Node, child: LN, *, visible: bool = True) -> None:
Expand Down
25 changes: 25 additions & 0 deletions tests/test_typed_comments.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,25 @@
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))
itxasos23 marked this conversation as resolved.
Show resolved Hide resolved


@pytest.mark.parametrize(
"comment_str",
[
"# Not a type comment: for sure.",
"# this is a comment about a :param called type",
"#type : the first space is mandatory",
],
)
def test_not_typed_comments(comment_str: str) -> None:
assert not is_type_comment(Leaf(token.COMMENT, comment_str))