Skip to content

Commit

Permalink
Fix is_type_comment function issue psf#2097
Browse files Browse the repository at this point in the history
  • Loading branch information
triet301 committed Feb 24, 2025
1 parent 256f342 commit 94b3915
Show file tree
Hide file tree
Showing 3 changed files with 51 additions and 2 deletions.
15 changes: 15 additions & 0 deletions src/black/lines.py
Original file line number Diff line number Diff line change
Expand Up @@ -18,6 +18,7 @@
is_multiline_string,
is_one_sequence_between,
is_type_comment,
normalize_type_comment,
is_type_ignore_comment,
is_with_or_async_with_stmt,
make_simple_prefix,
Expand Down Expand Up @@ -62,6 +63,11 @@ def append(
Inline comments are put aside.
"""

# Normalize type comments before processing
if leaf.type in {token.COMMENT, STANDALONE_COMMENT}:
normalize_type_comment(leaf)

has_value = (
leaf.type in BRACKETS
# empty fstring-middles must not be truncated
Expand Down Expand Up @@ -97,6 +103,10 @@ def append_safe(self, leaf: Leaf, preformatted: bool = False) -> None:
Raises ValueError when any `leaf` is appended after a standalone comment
or when a standalone comment is not the first leaf on the line.
"""
# Normalize type comments before appending
if leaf.type in {token.COMMENT, STANDALONE_COMMENT}:
normalize_type_comment(leaf)

if (
self.bracket_tracker.depth == 0
or self.bracket_tracker.any_open_for_or_lambda()
Expand Down Expand Up @@ -379,6 +389,11 @@ def has_magic_trailing_comma(self, closing: Leaf) -> bool:

def append_comment(self, comment: Leaf) -> bool:
"""Add an inline or standalone comment to the line."""

# Normalize type comments before deciding placement
if comment.type in {token.COMMENT, STANDALONE_COMMENT}:
normalize_type_comment(comment)

if (
comment.type == STANDALONE_COMMENT
and self.bracket_tracker.any_open_brackets()
Expand Down
17 changes: 15 additions & 2 deletions 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 collections.abc import Iterator
from typing import Final, Generic, Literal, Optional, TypeVar, Union
Expand Down Expand Up @@ -920,15 +921,27 @@ def is_async_stmt_or_funcdef(leaf: Leaf) -> bool:
)


def is_type_comment(leaf: Leaf) -> bool:
def is_type_comment(leaf: Leaf, suffix: str = "") -> bool:
"""Return True if the given leaf is a type comment. This function should only
be used for general type comments (excluding ignore annotations, which should
use `is_type_ignore_comment`). Note that general type comments are no longer
used in modern version of Python, this function may be deprecated in the future."""
t = leaf.type
v = leaf.value
return t in {token.COMMENT, STANDALONE_COMMENT} and v.startswith("# type:")
# Match "# type:" with optional spaces after "#" and before "type:"
pattern = r"^#\s*type\s*:" + re.escape(suffix)
return t in {token.COMMENT, STANDALONE_COMMENT} and re.match(pattern, v) is not None

def normalize_type_comment(leaf: Leaf) -> None:
"""Normalize the formatting of a type comment."""
if not is_type_comment(leaf):
return
# Extract the type annotation part after "# type:"
match = re.match(r"^#\s*type\s*:\s*(.+)$", leaf.value.strip())
if match:
annotation = match.group(1).strip() # Remove leading/trailing spaces from annotation
leaf.value = f"# type: {annotation}" # Standardize to "# type: <annotation>"


def is_type_ignore_comment(leaf: Leaf) -> bool:
"""Return True if the given leaf is a type comment with ignore annotation."""
Expand Down
21 changes: 21 additions & 0 deletions tests/data/cases/is_type_comment.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,21 @@
from black import format_str, FileMode
from black.nodes import is_type_comment, Leaf, STANDALONE_COMMENT

def test_is_type_comment_recognition():

assert is_type_comment(Leaf(STANDALONE_COMMENT, "# type: int"))
assert is_type_comment(Leaf(STANDALONE_COMMENT, "# type: int"))
assert is_type_comment(Leaf(STANDALONE_COMMENT, "# type: int"))
assert not is_type_comment(Leaf(STANDALONE_COMMENT, "# nope: int"))

def test_type_comment_normalization():

src = "x = 1 # type: int\n"
expected = "x = 1 # type: int\n"
result = format_str(src, mode=FileMode())
assert result == expected

src = "def foo(): # type: str\n pass"
expected = "def foo(): # type: str\n pass\n"
result = format_str(src, mode=FileMode())
assert result == expected

0 comments on commit 94b3915

Please sign in to comment.