Handle comments with annotations

PiperOrigin-RevId: 679548102

src/pyink/__init__.py

@@ -69,7 +69,13 @@
 from pyink.lines import EmptyLineTracker, LinesBlock
 from pyink.mode import FUTURE_FLAG_TO_FEATURE, Feature, VERSION_TO_FEATURES
 from pyink.mode import Mode as Mode  # re-exported
-from pyink.mode import Preview, QuoteStyle, TargetVersion, supports_feature
+from pyink.mode import (
+    DEFAULT_ANNOTATION_PRAGMAS,
+    Preview,
+    QuoteStyle,
+    TargetVersion,
+    supports_feature,
+)
 from pyink.nodes import STARS, is_number_token, is_simple_decorator_expression, syms
 from pyink.output import color_diff, diff, dump_to_file, err, ipynb_diff, out
 from pyink.parsing import (  # noqa F401
@@ -368,6 +374,18 @@ def validate_regex(
         " notebook."
     ),
 )
+@click.option(
+    "--pyink-annotation-pragmas",
+    type=str,
+    multiple=True,
+    help=(
+        "Pyink won't split too long lines if they contain a comment with any of"
+        " the given annotation pragmas because it doesn't know to which part of"
+        " the line the annotation applies. The default annotation pragmas are:"
+        f" {', '.join(DEFAULT_ANNOTATION_PRAGMAS)}."
+    ),
+    default=[],
+)
 @click.option(
     "--pyink-use-majority-quotes",
     is_flag=True,
@@ -580,6 +598,7 @@ def main(  # noqa: C901
     pyink: bool,
     pyink_indentation: str,
     pyink_ipynb_indentation: str,
+    pyink_annotation_pragmas: List[str],
     pyink_use_majority_quotes: bool,
     quiet: bool,
     verbose: bool,
@@ -681,6 +700,9 @@ def main(  # noqa: C901
         is_pyink=pyink,
         pyink_indentation=int(pyink_indentation),
         pyink_ipynb_indentation=int(pyink_ipynb_indentation),
+        pyink_annotation_pragmas=(
+            tuple(pyink_annotation_pragmas) or DEFAULT_ANNOTATION_PRAGMAS
+        ),
         quote_style=(
             QuoteStyle.MAJORITY if pyink_use_majority_quotes else QuoteStyle.DOUBLE
         ),

src/pyink/comments.py

@@ -3,6 +3,7 @@
 from functools import lru_cache
 from typing import Collection, Final, Iterator, List, Optional, Tuple, Union
 
+from pyink import ink
 from pyink.mode import Mode, Preview
 from pyink.nodes import (
     CLOSING_BRACKETS,
@@ -376,15 +377,15 @@ def children_contains_fmt_on(container: LN) -> bool:
     return False
 
 
-def contains_pragma_comment(comment_list: List[Leaf]) -> bool:
+def contains_pragma_comment(comment_list: List[Leaf], mode: Mode) -> bool:
     """
     Returns:
         True iff one of the comments in @comment_list is a pragma used by one
         of the more common static analysis tools for python (e.g. mypy, flake8,
         pylint).
     """
     for comment in comment_list:
-        if comment.value.startswith(("# type:", "# noqa", "# pylint:")):
+        if ink.comment_contains_pragma(comment.value, mode):
             return True
 
     return False

src/pyink/ink.py

@@ -17,7 +17,7 @@
 
 from blib2to3.pgen2.token import ASYNC, FSTRING_START, NEWLINE, STRING
 from blib2to3.pytree import type_repr
-from pyink.mode import Quote
+from pyink.mode import Mode, Quote
 from pyink.nodes import LN, Leaf, Node, STANDALONE_COMMENT, Visitor, syms
 from pyink.strings import STRING_PREFIX_CHARS
 
@@ -63,6 +63,28 @@ def majority_quote(node: LN) -> Quote:
     return Quote.DOUBLE
 
 
+def comment_contains_pragma(comment: str, mode: Mode) -> bool:
+    """Check if the given string contains one of the pragma forms.
+
+    A pragma form can appear at the beginning of a comment:
+      # pytype: disable=attribute-error
+    or somewhere in the middle:
+      # some comment # type: ignore # another comment
+    or the comments can even be separated by a semicolon:
+      # some comment; noqa: E111; another comment
+
+    Args:
+      comment: The comment to check.
+      mode: The mode that defines which pragma forms to check for.
+
+    Returns:
+      True if the comment contains one of the pragma forms.
+    """
+    joined_pragma_expression = "|".join(mode.pyink_annotation_pragmas)
+    pragma_regex = re.compile(rf"([#|;] ?(?:{joined_pragma_expression}))")
+    return pragma_regex.search(comment) is not None
+
+
 def get_code_start(src: str) -> str:
     """Provides the first line where the code starts.
 

src/pyink/linegen.py

@@ -14,6 +14,7 @@
 else:
     from typing import Final, Literal
 
+from pyink import ink
 from pyink.brackets import (
     COMMA_PRIORITY,
     DOT_PRIORITY,
@@ -61,7 +62,6 @@
     is_stub_body,
     is_stub_suite,
     is_tuple_containing_walrus,
-    is_type_ignore_comment_string,
     is_vararg,
     is_walrus_assignment,
     is_yield,
@@ -272,6 +272,7 @@ def visit_dictsetmaker(self, node: Node) -> Iterator[Line]:
                         maybe_make_parens_invisible_in_atom(
                             child,
                             parent=node,
+                            mode=self.mode,
                             remove_brackets_around_comma=False,
                         )
                     else:
@@ -292,6 +293,7 @@ def visit_funcdef(self, node: Node) -> Iterator[Line]:
                     if maybe_make_parens_invisible_in_atom(
                         child,
                         parent=node,
+                        mode=self.mode,
                         remove_brackets_around_comma=False,
                     ):
                         wrap_in_parentheses(node, child, visible=False)
@@ -315,8 +317,10 @@ def visit_match_case(self, node: Node) -> Iterator[Line]:
     def visit_suite(self, node: Node) -> Iterator[Line]:
         """Visit a suite."""
         if (
-            self.mode.is_pyi or not self.mode.is_pyink
-        ) and is_stub_suite(node, self.mode):
+            self.mode.is_pyi
+            or not self.mode.is_pyink
+            and is_stub_suite(node, self.mode)
+        ):
             yield from self.visit(node.children[2])
         else:
             yield from self.visit_default(node)
@@ -395,7 +399,7 @@ def visit_power(self, node: Node) -> Iterator[Line]:
             ):
                 wrap_in_parentheses(node, leaf)
 
-        remove_await_parens(node)
+        remove_await_parens(node, mode=self.mode)
 
         yield from self.visit_default(node)
 
@@ -442,7 +446,9 @@ def foo(a: int, b: float = 7): ...
         def foo(a: (int), b: (float) = 7): ...
         """
         assert len(node.children) == 3
-        if maybe_make_parens_invisible_in_atom(node.children[2], parent=node):
+        if maybe_make_parens_invisible_in_atom(
+            node.children[2], parent=node, mode=self.mode
+        ):
             wrap_in_parentheses(node, node.children[2], visible=False)
 
         yield from self.visit_default(node)
@@ -688,7 +694,7 @@ def transform_line(
 
     transformers: List[Transformer]
     if (
-        not line.contains_uncollapsable_type_comments()
+        not line.contains_uncollapsable_pragma_comments()
         and not line.should_split_rhs
         and not line.magic_trailing_comma
         and (
@@ -1451,15 +1457,17 @@ def normalize_invisible_parens(  # noqa: C901
                 if maybe_make_parens_invisible_in_atom(
                     child,
                     parent=node,
+                    mode=mode,
                     remove_brackets_around_comma=True,
                 ):
                     wrap_in_parentheses(node, child, visible=False)
             elif isinstance(child, Node) and node.type == syms.with_stmt:
-                remove_with_parens(child, node)
+                remove_with_parens(child, node, mode=mode)
             elif child.type == syms.atom:
                 if maybe_make_parens_invisible_in_atom(
                     child,
                     parent=node,
+                    mode=mode,
                 ):
                     wrap_in_parentheses(node, child, visible=False)
             elif is_one_tuple(child):
@@ -1511,7 +1519,7 @@ def _normalize_import_from(parent: Node, child: LN, index: int) -> None:
         parent.append_child(Leaf(token.RPAR, ""))
 
 
-def remove_await_parens(node: Node) -> None:
+def remove_await_parens(node: Node, mode: Mode) -> None:
     if node.children[0].type == token.AWAIT and len(node.children) > 1:
         if (
             node.children[1].type == syms.atom
@@ -1520,6 +1528,7 @@ def remove_await_parens(node: Node) -> None:
             if maybe_make_parens_invisible_in_atom(
                 node.children[1],
                 parent=node,
+                mode=mode,
                 remove_brackets_around_comma=True,
             ):
                 wrap_in_parentheses(node, node.children[1], visible=False)
@@ -1588,7 +1597,7 @@ def _maybe_wrap_cms_in_parens(
         node.insert_child(1, new_child)
 
 
-def remove_with_parens(node: Node, parent: Node) -> None:
+def remove_with_parens(node: Node, parent: Node, mode: Mode) -> None:
     """Recursively hide optional parens in `with` statements."""
     # Removing all unnecessary parentheses in with statements in one pass is a tad
     # complex as different variations of bracketed statements result in pretty
@@ -1610,21 +1619,23 @@ def remove_with_parens(node: Node, parent: Node) -> None:
         if maybe_make_parens_invisible_in_atom(
             node,
             parent=parent,
+            mode=mode,
             remove_brackets_around_comma=True,
         ):
             wrap_in_parentheses(parent, node, visible=False)
         if isinstance(node.children[1], Node):
-            remove_with_parens(node.children[1], node)
+            remove_with_parens(node.children[1], node, mode=mode)
     elif node.type == syms.testlist_gexp:
         for child in node.children:
             if isinstance(child, Node):
-                remove_with_parens(child, node)
+                remove_with_parens(child, node, mode=mode)
     elif node.type == syms.asexpr_test and not any(
         leaf.type == token.COLONEQUAL for leaf in node.leaves()
     ):
         if maybe_make_parens_invisible_in_atom(
             node.children[0],
             parent=node,
+            mode=mode,
             remove_brackets_around_comma=True,
         ):
             wrap_in_parentheses(node, node.children[0], visible=False)
@@ -1633,6 +1644,7 @@ def remove_with_parens(node: Node, parent: Node) -> None:
 def maybe_make_parens_invisible_in_atom(
     node: LN,
     parent: LN,
+    mode: Mode,
     remove_brackets_around_comma: bool = False,
 ) -> bool:
     """If it's safe, make the parens in the atom `node` invisible, recursively.
@@ -1682,7 +1694,7 @@ def maybe_make_parens_invisible_in_atom(
         if (
             # If the prefix of `middle` includes a type comment with
             # ignore annotation, then we do not remove the parentheses
-            not is_type_ignore_comment_string(middle.prefix.strip())
+            not ink.comment_contains_pragma(middle.prefix.strip(), mode)
         ):
             first.value = ""
             if first.prefix.strip():
@@ -1692,6 +1704,7 @@ def maybe_make_parens_invisible_in_atom(
         maybe_make_parens_invisible_in_atom(
             middle,
             parent=parent,
+            mode=mode,
             remove_brackets_around_comma=remove_brackets_around_comma,
         )
 
@@ -1835,7 +1848,7 @@ def run_transformer(
         or not line.bracket_tracker.invisible
         or any(bracket.value for bracket in line.bracket_tracker.invisible)
         or line.contains_multiline_strings()
-        or result[0].contains_uncollapsable_type_comments()
+        or result[0].contains_uncollapsable_pragma_comments()
         or result[0].contains_unsplittable_type_ignore()
         or is_line_short_enough(result[0], mode=mode)
         # If any leaves have no parents (which _can_ occur since