Exclude string type annotations from ESP (psf#3462)

yilei · web-flow · commit a44dc3d59eb4 · 2022-12-20T13:38:35.000-08:00
diff --git a/CHANGES.md b/CHANGES.md
@@ -26,6 +26,8 @@
 - Long values in dict literals are now wrapped in parentheses; correspondingly
   unnecessary parentheses around short values in dict literals are now removed; long
   string lambda values are now wrapped in parentheses (#3440)
+- Exclude string type annotations from improved string processing; fix crash when the
+  return type annotation is stringified and spans across multiple lines (#3462)
 
 ### Configuration
 
diff --git a/src/black/brackets.py b/src/black/brackets.py
@@ -80,9 +80,12 @@ def mark(self, leaf: Leaf) -> None:
         within brackets a given leaf is. 0 means there are no enclosing brackets
         that started on this line.
 
-        If a leaf is itself a closing bracket, it receives an `opening_bracket`
-        field that it forms a pair with. This is a one-directional link to
-        avoid reference cycles.
+        If a leaf is itself a closing bracket and there is a matching opening
+        bracket earlier, it receives an `opening_bracket` field with which it forms a
+        pair. This is a one-directional link to avoid reference cycles. Closing
+        bracket without opening happens on lines continued from previous
+        breaks, e.g. `) -> "ReturnType":` as part of a funcdef where we place
+        the return type annotation on its own line of the previous closing RPAR.
 
         If a leaf is a delimiter (a token on which Black can split the line if
         needed) and it's on depth 0, its `id()` is stored in the tracker's
@@ -91,6 +94,13 @@ def mark(self, leaf: Leaf) -> None:
         if leaf.type == token.COMMENT:
             return
 
+        if (
+            self.depth == 0
+            and leaf.type in CLOSING_BRACKETS
+            and (self.depth, leaf.type) not in self.bracket_match
+        ):
+            return
+
         self.maybe_decrement_after_for_loop_variable(leaf)
         self.maybe_decrement_after_lambda_arguments(leaf)
         if leaf.type in CLOSING_BRACKETS:
diff --git a/src/black/nodes.py b/src/black/nodes.py
@@ -848,3 +848,15 @@ def is_string_token(nl: NL) -> TypeGuard[Leaf]:
 
 def is_number_token(nl: NL) -> TypeGuard[Leaf]:
     return nl.type == token.NUMBER
+
+
+def is_part_of_annotation(leaf: Leaf) -> bool:
+    """Returns whether this leaf is part of type annotations."""
+    ancestor = leaf.parent
+    while ancestor is not None:
+        if ancestor.prev_sibling and ancestor.prev_sibling.type == token.RARROW:
+            return True
+        if ancestor.parent and ancestor.parent.type == syms.tname:
+            return True
+        ancestor = ancestor.parent
+    return False
diff --git a/src/black/trans.py b/src/black/trans.py
@@ -30,7 +30,6 @@
 
 from mypy_extensions import trait
 
-from black.brackets import BracketMatchError
 from black.comments import contains_pragma_comment
 from black.lines import Line, append_leaves
 from black.mode import Feature
@@ -41,6 +40,7 @@
     is_empty_lpar,
     is_empty_par,
     is_empty_rpar,
+    is_part_of_annotation,
     parent_type,
     replace_child,
     syms,
@@ -351,7 +351,7 @@ class StringMerger(StringTransformer, CustomSplitMapMixin):
 
     Requirements:
         (A) The line contains adjacent strings such that ALL of the validation checks
-        listed in StringMerger.__validate_msg(...)'s docstring pass.
+        listed in StringMerger._validate_msg(...)'s docstring pass.
             OR
         (B) The line contains a string which uses line continuation backslashes.
 
@@ -377,6 +377,8 @@ def do_match(self, line: Line) -> TMatchResult:
                 and is_valid_index(i + 1)
                 and LL[i + 1].type == token.STRING
             ):
+                if is_part_of_annotation(leaf):
+                    return TErr("String is part of type annotation.")
                 return Ok(i)
 
             if leaf.type == token.STRING and "\\\n" in leaf.value:
@@ -454,7 +456,7 @@ def _merge_string_group(self, line: Line, string_idx: int) -> TResult[Line]:
 
         Returns:
             Ok(new_line), if ALL of the validation checks found in
-            __validate_msg(...) pass.
+            _validate_msg(...) pass.
                 OR
             Err(CannotTransform), otherwise.
         """
@@ -608,7 +610,7 @@ def make_naked(string: str, string_prefix: str) -> str:
     def _validate_msg(line: Line, string_idx: int) -> TResult[None]:
         """Validate (M)erge (S)tring (G)roup
 
-        Transform-time string validation logic for __merge_string_group(...).
+        Transform-time string validation logic for _merge_string_group(...).
 
         Returns:
             * Ok(None), if ALL validation checks (listed below) pass.
@@ -622,6 +624,11 @@ def _validate_msg(line: Line, string_idx: int) -> TResult[None]:
                 - The set of all string prefixes in the string group is of
                   length greater than one and is not equal to {"", "f"}.
                 - The string group consists of raw strings.
+                - The string group is stringified type annotations. We don't want to
+                  process stringified type annotations since pyright doesn't support
+                  them spanning multiple string values. (NOTE: mypy, pytype, pyre do
+                  support them, so we can change if pyright also gains support in the
+                  future. See https://github.com/microsoft/pyright/issues/4359.)
         """
         # We first check for "inner" stand-alone comments (i.e. stand-alone
         # comments that have a string leaf before them AND after them).
@@ -812,13 +819,7 @@ def do_transform(self, line: Line, string_idx: int) -> Iterator[TResult[Line]]:
 
         new_line = line.clone()
         new_line.comments = line.comments.copy()
-        try:
-            append_leaves(new_line, line, LL[: string_idx - 1])
-        except BracketMatchError:
-            # HACK: I believe there is currently a bug somewhere in
-            # right_hand_split() that is causing brackets to not be tracked
-            # properly by a shared BracketTracker.
-            append_leaves(new_line, line, LL[: string_idx - 1], preformatted=True)
+        append_leaves(new_line, line, LL[: string_idx - 1])
 
         string_leaf = Leaf(token.STRING, LL[string_idx].value)
         LL[string_idx - 1].remove()
diff --git a/tests/data/preview/long_strings__type_annotations.py b/tests/data/preview/long_strings__type_annotations.py
@@ -0,0 +1,59 @@
+def func(
+    arg1,
+    arg2,
+) -> Set["this_is_a_very_long_module_name.AndAVeryLongClasName"
+         ".WithAVeryVeryVeryVeryVeryLongSubClassName"]:
+  pass
+
+
+def func(
+    argument: (
+        "VeryLongClassNameWithAwkwardGenericSubtype[int] |"
+        "VeryLongClassNameWithAwkwardGenericSubtype[str]"
+    ),
+) -> (
+    "VeryLongClassNameWithAwkwardGenericSubtype[int] |"
+    "VeryLongClassNameWithAwkwardGenericSubtype[str]"
+):
+  pass
+
+
+def func(
+    argument: (
+        "int |"
+        "str"
+    ),
+) -> Set["int |"
+         " str"]:
+  pass
+
+
+# output
+
+
+def func(
+    arg1,
+    arg2,
+) -> Set[
+    "this_is_a_very_long_module_name.AndAVeryLongClasName"
+    ".WithAVeryVeryVeryVeryVeryLongSubClassName"
+]:
+    pass
+
+
+def func(
+    argument: (
+        "VeryLongClassNameWithAwkwardGenericSubtype[int] |"
+        "VeryLongClassNameWithAwkwardGenericSubtype[str]"
+    ),
+) -> (
+    "VeryLongClassNameWithAwkwardGenericSubtype[int] |"
+    "VeryLongClassNameWithAwkwardGenericSubtype[str]"
+):
+    pass
+
+
+def func(
+    argument: ("int |" "str"),
+) -> Set["int |" " str"]:
+    pass
