Simplify power operator hugging implementation

CHANGES.md

@@ -17,6 +17,8 @@
 
 <!-- Changes that affect Black's preview style -->
 
+- Simplify implementation of the power operator "hugging" logic (#4918)
+
 ### Configuration
 
 <!-- Changes to how Black can be configured -->

docs/the_black_code_style/future_style.md

@@ -15,6 +15,9 @@ Currently, the following features are included in the preview style:
 
 - `wrap_comprehension_in`: Wrap the `in` clause of list and dictionary comprehensions
   across lines if it would otherwise exceed the maximum line length.
+- `simplify_power_operator_hugging`: Use a simpler implementation of the power operator
+  "hugging" logic (removing whitespace around `**` in simple expressions), which applies
+  also in the rare case the exponentiation is split into separate lines.
 - `wrap_long_dict_values_in_parens`: Add parentheses around long values in dictionaries.
   ([see below](labels/wrap-long-dict-values))
 

src/black/linegen.py

@@ -679,6 +679,7 @@ def __post_init__(self) -> None:
         self.visit_guard = partial(v, keywords=Ø, parens={"if"})
 
 
+# Remove when `simplify_power_operator_hugging` becomes stable.
 def _hugging_power_ops_line_to_string(
     line: Line,
     features: Collection[Feature],
@@ -705,11 +706,15 @@ def transform_line(
 
     line_str = line_to_string(line)
 
-    # We need the line string when power operators are hugging to determine if we should
-    # split the line. Default to line_str, if no power operator are present on the line.
-    line_str_hugging_power_ops = (
-        _hugging_power_ops_line_to_string(line, features, mode) or line_str
-    )
+    if Preview.simplify_power_operator_hugging in mode:
+        line_str_hugging_power_ops = line_str
+    else:
+        # We need the line string when power operators are hugging to determine if we
+        # should split the line. Default to line_str, if no power operator are present
+        # on the line.
+        line_str_hugging_power_ops = (
+            _hugging_power_ops_line_to_string(line, features, mode) or line_str
+        )
 
     ll = mode.line_length
     sn = mode.string_normalization
@@ -794,9 +799,11 @@ def _rhs(
                 transformers = [delimiter_split, standalone_comment_split, rhs]
             else:
                 transformers = [rhs]
-    # It's always safe to attempt hugging of power operations and pretty much every line
-    # could match.
-    transformers.append(hug_power_op)
+
+    if Preview.simplify_power_operator_hugging not in mode:
+        # It's always safe to attempt hugging of power operations and pretty much every
+        # line could match.
+        transformers.append(hug_power_op)
 
     for transform in transformers:
         # We are accumulating lines in `result` because we might want to abort

src/black/mode.py

@@ -227,6 +227,7 @@ class Preview(Enum):
     string_processing = auto()
     hug_parens_with_braces_and_square_brackets = auto()
     wrap_comprehension_in = auto()
+    simplify_power_operator_hugging = auto()
     wrap_long_dict_values_in_parens = auto()
 
 

src/black/nodes.py

@@ -8,7 +8,7 @@
 from mypy_extensions import mypyc_attr
 
 from black.cache import CACHE_DIR
-from black.mode import Mode
+from black.mode import Mode, Preview
 from black.strings import get_string_prefix, has_triple_quotes
 from blib2to3 import pygram
 from blib2to3.pgen2 import token
@@ -416,6 +416,15 @@ def whitespace(leaf: Leaf, *, complex_subscript: bool, mode: Mode) -> str:
         if t == token.STAR:
             return NO
 
+    if Preview.simplify_power_operator_hugging in mode:
+        # Power operator hugging
+        if t == token.DOUBLESTAR and is_simple_exponentiation(p):
+            return NO
+        prevp = preceding_leaf(leaf)
+        if prevp and prevp.type == token.DOUBLESTAR:
+            if prevp.parent and is_simple_exponentiation(prevp.parent):
+                return NO
+
     return SPACE
 
 
@@ -543,6 +552,25 @@ def is_arith_like(node: LN) -> bool:
     }
 
 
+def is_simple_exponentiation(node: LN) -> bool:
+    """Whether whitespace around `**` should be removed."""
+
+    def is_simple(node: LN) -> bool:
+        if isinstance(node, Leaf):
+            return node.type in (token.NAME, token.NUMBER, token.DOT, token.DOUBLESTAR)
+        elif node.type == syms.factor:  # unary operators
+            return is_simple(node.children[1])
+        else:
+            return all(is_simple(child) for child in node.children)
+
+    return (
+        node.type == syms.power
+        and len(node.children) >= 3
+        and node.children[-2].type == token.DOUBLESTAR
+        and is_simple(node)
+    )
+
+
 def is_docstring(node: NL) -> bool:
     if isinstance(node, Leaf):
         if node.type != token.STRING:

src/black/resources/black.schema.json

@@ -83,6 +83,7 @@
           "string_processing",
           "hug_parens_with_braces_and_square_brackets",
           "wrap_comprehension_in",
+          "simplify_power_operator_hugging",
           "wrap_long_dict_values_in_parens"
         ]
       },

src/black/trans.py

@@ -66,6 +66,7 @@ def TErr(err_msg: str) -> Err[CannotTransform]:
     return Err(cant_transform)
 
 
+# Remove when `simplify_power_operator_hugging` becomes stable.
 def hug_power_op(
     line: Line, features: Collection[Feature], mode: Mode
 ) -> Iterator[Line]:
@@ -133,6 +134,7 @@ def is_simple_operand(index: int, kind: Literal[1, -1]) -> bool:
     yield new_line
 
 
+# Remove when `simplify_power_operator_hugging` becomes stable.
 def handle_is_simple_look_up_prev(line: Line, index: int, disallowed: set[int]) -> bool:
     """
     Handling the determination of is_simple_lookup for the lines prior to the doublestar
@@ -155,6 +157,7 @@ def handle_is_simple_look_up_prev(line: Line, index: int, disallowed: set[int])
     return True
 
 
+# Remove when `simplify_power_operator_hugging` becomes stable.
 def handle_is_simple_lookup_forward(
     line: Line, index: int, disallowed: set[int]
 ) -> bool:
@@ -181,6 +184,7 @@ def handle_is_simple_lookup_forward(
     return True
 
 
+# Remove when `simplify_power_operator_hugging` becomes stable.
 def is_expression_chained(chained_leaves: list[Leaf]) -> bool:
     """
     Function to determine if the variable is a chained call.

tests/data/cases/preview_simplify_power_operator_hugging.py

@@ -0,0 +1,153 @@
+# flags: --preview
+# This is a copy of `power_op_spacing.py`. Remove when `simplify_power_operator_hugging` becomes stable.
+
+def function(**kwargs):
+    t = a**2 + b**3
+    return t ** 2
+
+
+def function_replace_spaces(**kwargs):
+    t = a **2 + b** 3 + c ** 4
+
+
+def function_dont_replace_spaces():
+    {**a, **b, **c}
+
+
+a = 5**~4
+b = 5 ** f()
+c = -(5**2)
+d = 5 ** f["hi"]
+e = lazy(lambda **kwargs: 5)
+f = f() ** 5
+g = a.b**c.d
+h = 5 ** funcs.f()
+i = funcs.f() ** 5
+j = super().name ** 5
+k = [(2**idx, value) for idx, value in pairs]
+l = mod.weights_[0] == pytest.approx(0.95**100, abs=0.001)
+m = [([2**63], [1, 2**63])]
+n = count <= 10**5
+o = settings(max_examples=10**6)
+p = {(k, k**2): v**2 for k, v in pairs}
+q = [10**i for i in range(6)]
+r = x**y
+s = 1 ** 1
+t = (
+    1
+    ** 1
+    **1
+    ** 1
+)
+
+a = 5.0**~4.0
+b = 5.0 ** f()
+c = -(5.0**2.0)
+d = 5.0 ** f["hi"]
+e = lazy(lambda **kwargs: 5)
+f = f() ** 5.0
+g = a.b**c.d
+h = 5.0 ** funcs.f()
+i = funcs.f() ** 5.0
+j = super().name ** 5.0
+k = [(2.0**idx, value) for idx, value in pairs]
+l = mod.weights_[0] == pytest.approx(0.95**100, abs=0.001)
+m = [([2.0**63.0], [1.0, 2**63.0])]
+n = count <= 10**5.0
+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)]
+s = 1.0 ** 1.0
+t = (
+    1.0
+    ** 1.0
+    **1.0
+    ** 1.0
+)
+
+
+# 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
+# This is a copy of `power_op_spacing.py`. Remove when `simplify_power_operator_hugging` becomes stable.
+
+
+def function(**kwargs):
+    t = a**2 + b**3
+    return t**2
+
+
+def function_replace_spaces(**kwargs):
+    t = a**2 + b**3 + c**4
+
+
+def function_dont_replace_spaces():
+    {**a, **b, **c}
+
+
+a = 5**~4
+b = 5 ** f()
+c = -(5**2)
+d = 5 ** f["hi"]
+e = lazy(lambda **kwargs: 5)
+f = f() ** 5
+g = a.b**c.d
+h = 5 ** funcs.f()
+i = funcs.f() ** 5
+j = super().name ** 5
+k = [(2**idx, value) for idx, value in pairs]
+l = mod.weights_[0] == pytest.approx(0.95**100, abs=0.001)
+m = [([2**63], [1, 2**63])]
+n = count <= 10**5
+o = settings(max_examples=10**6)
+p = {(k, k**2): v**2 for k, v in pairs}
+q = [10**i for i in range(6)]
+r = x**y
+s = 1**1
+t = 1**1**1**1
+
+a = 5.0**~4.0
+b = 5.0 ** f()
+c = -(5.0**2.0)
+d = 5.0 ** f["hi"]
+e = lazy(lambda **kwargs: 5)
+f = f() ** 5.0
+g = a.b**c.d
+h = 5.0 ** funcs.f()
+i = funcs.f() ** 5.0
+j = super().name ** 5.0
+k = [(2.0**idx, value) for idx, value in pairs]
+l = mod.weights_[0] == pytest.approx(0.95**100, abs=0.001)
+m = [([2.0**63.0], [1.0, 2**63.0])]
+n = count <= 10**5.0
+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)]
+s = 1.0**1.0
+t = 1.0**1.0**1.0**1.0
+
+
+# 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
+)