Copy over comments when hugging power ops (#2874)

Otherwise they'd be deleted which was a regression in 22.1.0 (oops! my
bad!). Also type comments are now tracked in the AST safety check on all
compatible platforms to error out if this happens again.

Overall the line rewriting code has been rewritten to do "the right
thing (tm)", I hope this fixes other potential bugs in the code (fwiw I
got to drop the bugfix in blib2to3.pytree.Leaf.clone since now bracket
metadata is properly copied over).

Fixes #2873

CHANGES.md

@@ -10,6 +10,9 @@
 
 <!-- Changes that affect Black's stable style -->
 
+- Comments are no longer deleted when a line had spaces removed around power operators
+  (#2874)
+
 ### Preview style
 
 <!-- Changes that affect Black's preview style -->
@@ -47,6 +50,10 @@
 
 <!-- Changes to the parser or to version autodetection -->
 
+- Type comments are now included in the AST equivalence check consistently so accidental
+  deletion raises an error. Though type comments can't be tracked when running on PyPy
+  3.7 due to standard library limitations. (#2874)
+
 ### Performance
 
 <!-- Changes that improve Black's performance. -->

src/black/parsing.py

@@ -152,14 +152,22 @@ def parse_single_version(
     src: str, version: Tuple[int, int]
 ) -> Union[ast.AST, ast3.AST]:
     filename = "<unknown>"
-    # typed_ast is needed because of feature version limitations in the builtin ast
+    # typed-ast is needed because of feature version limitations in the builtin ast 3.8>
     if sys.version_info >= (3, 8) and version >= (3,):
-        return ast.parse(src, filename, feature_version=version)
-    elif version >= (3,):
-        if _IS_PYPY:
-            return ast3.parse(src, filename)
+        return ast.parse(src, filename, feature_version=version, type_comments=True)
+
+    if _IS_PYPY:
+        # PyPy 3.7 doesn't support type comment tracking which is not ideal, but there's
+        # not much we can do as typed-ast won't work either.
+        if sys.version_info >= (3, 8):
+            return ast3.parse(src, filename, type_comments=True)
         else:
-            return ast3.parse(src, filename, feature_version=version[1])
+            return ast3.parse(src, filename)
+    else:
+        # Typed-ast is guaranteed to be used here and automatically tracks type
+        # comments separately.
+        return ast3.parse(src, filename, feature_version=version[1])
+
     raise AssertionError("INTERNAL ERROR: Tried parsing unsupported Python version!")
 
 

src/black/trans.py

@@ -121,7 +121,7 @@ def is_simple_operand(index: int, kind: Literal["base", "exponent"]) -> bool:
 
         return False
 
-    leaves: List[Leaf] = []
+    new_line = line.clone()
     should_hug = False
     for idx, leaf in enumerate(line.leaves):
         new_leaf = leaf.clone()
@@ -139,18 +139,14 @@ def is_simple_operand(index: int, kind: Literal["base", "exponent"]) -> bool:
         if should_hug:
             new_leaf.prefix = ""
 
-        leaves.append(new_leaf)
+        # We have to be careful to make a new line properly:
+        # - bracket related metadata must be maintained (handled by Line.append)
+        # - comments need to copied over, updating the leaf IDs they're attached to
+        new_line.append(new_leaf, preformatted=True)
+        for comment_leaf in line.comments_after(leaf):
+            new_line.append(comment_leaf, preformatted=True)
 
-    yield Line(
-        mode=line.mode,
-        depth=line.depth,
-        leaves=leaves,
-        comments=line.comments,
-        bracket_tracker=line.bracket_tracker,
-        inside_brackets=line.inside_brackets,
-        should_split_rhs=line.should_split_rhs,
-        magic_trailing_comma=line.magic_trailing_comma,
-    )
+    yield new_line
 
 
 class StringTransformer(ABC):

src/blib2to3/pytree.py

@@ -451,7 +451,6 @@ def clone(self) -> "Leaf":
             self.value,
             (self.prefix, (self.lineno, self.column)),
             fixers_applied=self.fixers_applied,
-            opening_bracket=self.opening_bracket,
         )
 
     def leaves(self) -> Iterator["Leaf"]:

tests/data/simple_cases/power_op_spacing.py

@@ -49,6 +49,20 @@ def function_dont_replace_spaces():
 q = [10.5**i for i in range(6)]
 
 
+# WE SHOULD DEFINITELY NOT EAT THESE COMMENTS (https://github.com/psf/black/issues/2873)
+if hasattr(view, "sum_of_weights"):
+    return np.divide(  # type: ignore[no-any-return]
+        view.variance,  # type: ignore[union-attr]
+        view.sum_of_weights,  # type: ignore[union-attr]
+        out=np.full(view.sum_of_weights.shape, np.nan),  # type: ignore[union-attr]
+        where=view.sum_of_weights**2 > view.sum_of_weights_squared,  # type: ignore[union-attr]
+    )
+
+return np.divide(
+    where=view.sum_of_weights_of_weight_long**2 > view.sum_of_weights_squared,  # type: ignore
+)
+
+
 # output
 
 
@@ -101,3 +115,17 @@ def function_dont_replace_spaces():
 o = settings(max_examples=10**6.0)
 p = {(k, k**2): v**2.0 for k, v in pairs}
 q = [10.5**i for i in range(6)]
+
+
+# WE SHOULD DEFINITELY NOT EAT THESE COMMENTS (https://github.com/psf/black/issues/2873)
+if hasattr(view, "sum_of_weights"):
+    return np.divide(  # type: ignore[no-any-return]
+        view.variance,  # type: ignore[union-attr]
+        view.sum_of_weights,  # type: ignore[union-attr]
+        out=np.full(view.sum_of_weights.shape, np.nan),  # type: ignore[union-attr]
+        where=view.sum_of_weights**2 > view.sum_of_weights_squared,  # type: ignore[union-attr]
+    )
+
+return np.divide(
+    where=view.sum_of_weights_of_weight_long**2 > view.sum_of_weights_squared,  # type: ignore
+)