Fix comments getting removed from inside parenthesized strings (#3909)

[etc/vim.git] / src / black / linegen.py

diff --git a/src/black/linegen.py b/src/black/linegen.py
index 95d5583c5f53fd1ab97287c78a41cc90d3d8c33d..9ddd4619f691275a8a1ce9b70d4bce7348596281 100644 (file)
--- a/src/black/linegen.py
+++ b/src/black/linegen.py
@@ -2,7 +2,7 @@
 Generating lines of code.
 """
 import sys
 Generating lines of code.
 """
 import sys

-from dataclasses import dataclass, replace
+from dataclasses import replace

 from enum import Enum, auto
 from functools import partial, wraps
 from typing import Collection, Iterator, List, Optional, Set, Union, cast
 from enum import Enum, auto
 from functools import partial, wraps
 from typing import Collection, Iterator, List, Optional, Set, Union, cast


@@ -16,6 +16,7 @@ from black.brackets import (
 from black.comments import FMT_OFF, generate_comments, list_comments
 from black.lines import (
     Line,
 from black.comments import FMT_OFF, generate_comments, list_comments
 from black.lines import (
     Line,

+    RHSResult,

     append_leaves,
     can_be_split,
     can_omit_invisible_parens,
     append_leaves,
     can_be_split,
     can_omit_invisible_parens,


@@ -35,6 +36,7 @@ from black.nodes import (
     Visitor,
     ensure_visible,
     is_arith_like,
     Visitor,
     ensure_visible,
     is_arith_like,

+    is_async_stmt_or_funcdef,

     is_atom_with_invisible_parens,
     is_docstring,
     is_empty_tuple,
     is_atom_with_invisible_parens,
     is_docstring,
     is_empty_tuple,


@@ -47,6 +49,7 @@ from black.nodes import (
     is_stub_body,
     is_stub_suite,
     is_tuple_containing_walrus,
     is_stub_body,
     is_stub_suite,
     is_tuple_containing_walrus,

+    is_type_ignore_comment_string,

     is_vararg,
     is_walrus_assignment,
     is_yield,
     is_vararg,
     is_walrus_assignment,
     is_yield,


@@ -109,6 +112,17 @@ class LineGenerator(Visitor[Line]):
             self.current_line.depth += indent
             return  # Line is empty, don't emit. Creating a new one unnecessary.
 
             self.current_line.depth += indent
             return  # Line is empty, don't emit. Creating a new one unnecessary.
 

+        if (
+            Preview.improved_async_statements_handling in self.mode
+            and len(self.current_line.leaves) == 1
+            and is_async_stmt_or_funcdef(self.current_line.leaves[0])
+        ):
+            # Special case for async def/for/with statements. `visit_async_stmt`
+            # adds an `ASYNC` leaf then visits the child def/for/with statement
+            # nodes. Line yields from those nodes shouldn't treat the former
+            # `ASYNC` leaf as a complete line.
+            return
+

         complete_line = self.current_line
         self.current_line = Line(mode=self.mode, depth=complete_line.depth + indent)
         yield complete_line
         complete_line = self.current_line
         self.current_line = Line(mode=self.mode, depth=complete_line.depth + indent)
         yield complete_line


@@ -202,6 +216,18 @@ class LineGenerator(Visitor[Line]):
 
             yield from self.visit(child)
 
 
             yield from self.visit(child)
 

+    def visit_typeparams(self, node: Node) -> Iterator[Line]:
+        yield from self.visit_default(node)
+        node.children[0].prefix = ""
+
+    def visit_typevartuple(self, node: Node) -> Iterator[Line]:
+        yield from self.visit_default(node)
+        node.children[1].prefix = ""
+
+    def visit_paramspec(self, node: Node) -> Iterator[Line]:
+        yield from self.visit_default(node)
+        node.children[1].prefix = ""
+

     def visit_dictsetmaker(self, node: Node) -> Iterator[Line]:
         if Preview.wrap_long_dict_values_in_parens in self.mode:
             for i, child in enumerate(node.children):
     def visit_dictsetmaker(self, node: Node) -> Iterator[Line]:
         if Preview.wrap_long_dict_values_in_parens in self.mode:
             for i, child in enumerate(node.children):


@@ -255,7 +281,9 @@ class LineGenerator(Visitor[Line]):
 
     def visit_suite(self, node: Node) -> Iterator[Line]:
         """Visit a suite."""
 
     def visit_suite(self, node: Node) -> Iterator[Line]:
         """Visit a suite."""

-        if self.mode.is_pyi and is_stub_suite(node):
+        if (
+            self.mode.is_pyi or Preview.dummy_implementations in self.mode
+        ) and is_stub_suite(node):

             yield from self.visit(node.children[2])
         else:
             yield from self.visit_default(node)
             yield from self.visit(node.children[2])
         else:
             yield from self.visit_default(node)


@@ -270,7 +298,9 @@ class LineGenerator(Visitor[Line]):
 
         is_suite_like = node.parent and node.parent.type in STATEMENT
         if is_suite_like:
 
         is_suite_like = node.parent and node.parent.type in STATEMENT
         if is_suite_like:

-            if self.mode.is_pyi and is_stub_body(node):
+            if (
+                self.mode.is_pyi or Preview.dummy_implementations in self.mode
+            ) and is_stub_body(node):

                 yield from self.visit_default(node)
             else:
                 yield from self.line(+1)
                 yield from self.visit_default(node)
             else:
                 yield from self.line(+1)


@@ -279,7 +309,7 @@ class LineGenerator(Visitor[Line]):
 
         else:
             if (
 
         else:
             if (

-                not self.mode.is_pyi
+                not (self.mode.is_pyi or Preview.dummy_implementations in self.mode)

                 or not node.parent
                 or not is_stub_suite(node.parent)
             ):
                 or not node.parent
                 or not is_stub_suite(node.parent)
             ):


@@ -300,8 +330,11 @@ class LineGenerator(Visitor[Line]):
                 break
 
         internal_stmt = next(children)
                 break
 
         internal_stmt = next(children)

-        for child in internal_stmt.children:
-            yield from self.visit(child)
+        if Preview.improved_async_statements_handling in self.mode:
+            yield from self.visit(internal_stmt)
+        else:
+            for child in internal_stmt.children:
+                yield from self.visit(child)

 
     def visit_decorators(self, node: Node) -> Iterator[Line]:
         """Visit decorators."""
 
     def visit_decorators(self, node: Node) -> Iterator[Line]:
         """Visit decorators."""


@@ -364,6 +397,24 @@ class LineGenerator(Visitor[Line]):
             node.insert_child(index, Node(syms.atom, [lpar, operand, rpar]))
         yield from self.visit_default(node)
 
             node.insert_child(index, Node(syms.atom, [lpar, operand, rpar]))
         yield from self.visit_default(node)
 

+    def visit_tname(self, node: Node) -> Iterator[Line]:
+        """
+        Add potential parentheses around types in function parameter lists to be made
+        into real parentheses in case the type hint is too long to fit on a line
+        Examples:
+        def foo(a: int, b: float = 7): ...
+
+        ->
+
+        def foo(a: (int), b: (float) = 7): ...
+        """
+        if Preview.parenthesize_long_type_hints in self.mode:
+            assert len(node.children) == 3
+            if maybe_make_parens_invisible_in_atom(node.children[2], parent=node):
+                wrap_in_parentheses(node, node.children[2], visible=False)
+
+        yield from self.visit_default(node)
+

     def visit_STRING(self, leaf: Leaf) -> Iterator[Line]:
         if Preview.hex_codes_in_unicode_sequences in self.mode:
             normalize_unicode_escape_sequences(leaf)
     def visit_STRING(self, leaf: Leaf) -> Iterator[Line]:
         if Preview.hex_codes_in_unicode_sequences in self.mode:
             normalize_unicode_escape_sequences(leaf)


@@ -465,7 +516,14 @@ class LineGenerator(Visitor[Line]):
         self.visit_except_clause = partial(v, keywords={"except"}, parens={"except"})
         self.visit_with_stmt = partial(v, keywords={"with"}, parens={"with"})
         self.visit_classdef = partial(v, keywords={"class"}, parens=Ø)
         self.visit_except_clause = partial(v, keywords={"except"}, parens={"except"})
         self.visit_with_stmt = partial(v, keywords={"with"}, parens={"with"})
         self.visit_classdef = partial(v, keywords={"class"}, parens=Ø)

-        self.visit_expr_stmt = partial(v, keywords=Ø, parens=ASSIGNMENTS)
+
+        # When this is moved out of preview, add ":" directly to ASSIGNMENTS in nodes.py
+        if Preview.parenthesize_long_type_hints in self.mode:
+            assignments = ASSIGNMENTS | {":"}
+        else:
+            assignments = ASSIGNMENTS
+        self.visit_expr_stmt = partial(v, keywords=Ø, parens=assignments)
+

         self.visit_return_stmt = partial(v, keywords={"return"}, parens={"return"})
         self.visit_import_from = partial(v, keywords=Ø, parens={"import"})
         self.visit_del_stmt = partial(v, keywords=Ø, parens={"del"})
         self.visit_return_stmt = partial(v, keywords={"return"}, parens={"return"})
         self.visit_import_from = partial(v, keywords=Ø, parens={"import"})
         self.visit_del_stmt = partial(v, keywords=Ø, parens={"del"})


@@ -647,17 +705,6 @@ def left_hand_split(
             yield result
 
 
             yield result
 
 

-@dataclass
-class _RHSResult:
-    """Intermediate split result from a right hand split."""
-
-    head: Line
-    body: Line
-    tail: Line
-    opening_bracket: Leaf
-    closing_bracket: Leaf
-
-

 def right_hand_split(
     line: Line,
     mode: Mode,
 def right_hand_split(
     line: Line,
     mode: Mode,


@@ -681,7 +728,7 @@ def right_hand_split(
 def _first_right_hand_split(
     line: Line,
     omit: Collection[LeafID] = (),
 def _first_right_hand_split(
     line: Line,
     omit: Collection[LeafID] = (),

-) -> _RHSResult:
+) -> RHSResult:

     """Split the line into head, body, tail starting with the last bracket pair.
 
     Note: this function should not have side effects. It's relied upon by
     """Split the line into head, body, tail starting with the last bracket pair.
 
     Note: this function should not have side effects. It's relied upon by


@@ -723,11 +770,11 @@ def _first_right_hand_split(
         tail_leaves, line, opening_bracket, component=_BracketSplitComponent.tail
     )
     bracket_split_succeeded_or_raise(head, body, tail)
         tail_leaves, line, opening_bracket, component=_BracketSplitComponent.tail
     )
     bracket_split_succeeded_or_raise(head, body, tail)

-    return _RHSResult(head, body, tail, opening_bracket, closing_bracket)
+    return RHSResult(head, body, tail, opening_bracket, closing_bracket)

 
 
 def _maybe_split_omitting_optional_parens(
 
 
 def _maybe_split_omitting_optional_parens(

-    rhs: _RHSResult,
+    rhs: RHSResult,

     line: Line,
     mode: Mode,
     features: Collection[Feature] = (),
     line: Line,
     mode: Mode,
     features: Collection[Feature] = (),


@@ -747,11 +794,11 @@ def _maybe_split_omitting_optional_parens(
         # there are no standalone comments in the body
         and not rhs.body.contains_standalone_comments(0)
         # and we can actually remove the parens
         # there are no standalone comments in the body
         and not rhs.body.contains_standalone_comments(0)
         # and we can actually remove the parens

-        and can_omit_invisible_parens(rhs.body, mode.line_length)
+        and can_omit_invisible_parens(rhs, mode.line_length)

     ):
         omit = {id(rhs.closing_bracket), *omit}
         try:
     ):
         omit = {id(rhs.closing_bracket), *omit}
         try:

-            # The _RHSResult Omitting Optional Parens.
+            # The RHSResult Omitting Optional Parens.

             rhs_oop = _first_right_hand_split(line, omit=omit)
             if not (
                 Preview.prefer_splitting_right_hand_side_of_assignments in line.mode
             rhs_oop = _first_right_hand_split(line, omit=omit)
             if not (
                 Preview.prefer_splitting_right_hand_side_of_assignments in line.mode


@@ -803,7 +850,7 @@ def _maybe_split_omitting_optional_parens(
             yield result
 
 
             yield result
 
 

-def _prefer_split_rhs_oop(rhs_oop: _RHSResult, mode: Mode) -> bool:
+def _prefer_split_rhs_oop(rhs_oop: RHSResult, mode: Mode) -> bool:

     """
     Returns whether we should prefer the result from a split omitting optional parens.
     """
     """
     Returns whether we should prefer the result from a split omitting optional parens.
     """


@@ -901,6 +948,13 @@ def bracket_split_build_line(
                     )
                     if isinstance(node, Node) and isinstance(node.prev_sibling, Leaf)
                 )
                     )
                     if isinstance(node, Node) and isinstance(node.prev_sibling, Leaf)
                 )

+                # Except the false negatives above for PEP 604 unions where we
+                # can't add the comma.
+                and not (
+                    leaves[0].parent
+                    and leaves[0].parent.next_sibling
+                    and leaves[0].parent.next_sibling.type == token.VBAR
+                )

             )
 
             if original.is_import or no_commas:
             )
 
             if original.is_import or no_commas:


@@ -1339,7 +1393,7 @@ def maybe_make_parens_invisible_in_atom(
     Returns whether the node should itself be wrapped in invisible parentheses.
     """
     if (
     Returns whether the node should itself be wrapped in invisible parentheses.
     """
     if (

-        node.type != syms.atom
+        node.type not in (syms.atom, syms.expr)

         or is_empty_tuple(node)
         or is_one_tuple(node)
         or (is_yield(node) and parent.type != syms.expr_stmt)
         or is_empty_tuple(node)
         or is_one_tuple(node)
         or (is_yield(node) and parent.type != syms.expr_stmt)


@@ -1363,6 +1417,7 @@ def maybe_make_parens_invisible_in_atom(
             syms.except_clause,
             syms.funcdef,
             syms.with_stmt,
             syms.except_clause,
             syms.funcdef,
             syms.with_stmt,

+            syms.tname,

             # these ones aren't useful to end users, but they do please fuzzers
             syms.for_stmt,
             syms.del_stmt,
             # these ones aren't useful to end users, but they do please fuzzers
             syms.for_stmt,
             syms.del_stmt,


@@ -1375,8 +1430,13 @@ def maybe_make_parens_invisible_in_atom(
     if is_lpar_token(first) and is_rpar_token(last):
         middle = node.children[1]
         # make parentheses invisible
     if is_lpar_token(first) and is_rpar_token(last):
         middle = node.children[1]
         # make parentheses invisible

-        first.value = ""
-        last.value = ""
+        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())
+        ):
+            first.value = ""
+            last.value = ""

         maybe_make_parens_invisible_in_atom(
             middle,
             parent=parent,
         maybe_make_parens_invisible_in_atom(
             middle,
             parent=parent,