Hug power operators if its operands are "simple" (#2726)

Since power operators almost always have the highest binding power in expressions, it's often more readable to hug it with its operands. The main exception to this is when its operands are non-trivial in which case the power operator will not hug, the rule for this is the following:

> For power ops, an operand is considered "simple" if it's only a NAME, numeric CONSTANT, or attribute access (chained attribute access is allowed), with or without a preceding unary operator. 

Fixes GH-538.
Closes GH-2095.

diff-shades results: https://gist.github.com/ichard26/ca6c6ad4bd1de5152d95418c8645354b

Co-authored-by: Diego <dpalma@evernote.com>
Co-authored-by: Felix Hildén <felix.hilden@gmail.com>
Co-authored-by: Jelle Zijlstra <jelle.zijlstra@gmail.com>

CHANGES.md

@@ -20,6 +20,7 @@
 - Tuple unpacking on `return` and `yield` constructs now implies 3.8+ (#2700)
 - Unparenthesized tuples on annotated assignments (e.g
   `values: Tuple[int, ...] = 1, 2, 3`) now implies 3.8+ (#2708)
+- Remove spaces around power operators if both operands are simple (#2726)
 - Allow setting custom cache directory on all platforms with environment variable
   `BLACK_CACHE_DIR` (#2739).
 - Text coloring added in the final statistics (#2712)

docs/the_black_code_style/current_style.md

@@ -284,6 +284,26 @@ multiple lines. This is so that _Black_ is compliant with the recent changes in
 [PEP 8](https://www.python.org/dev/peps/pep-0008/#should-a-line-break-before-or-after-a-binary-operator)
 style guide, which emphasizes that this approach improves readability.
 
+Almost all operators will be surrounded by single spaces, the only exceptions are unary
+operators (`+`, `-`, and `~`), and power operators when both operands are simple. For
+powers, an operand is considered simple if it's only a NAME, numeric CONSTANT, or
+attribute access (chained attribute access is allowed), with or without a preceding
+unary operator.
+
+```python
+# For example, these won't be surrounded by whitespace
+a = x**y
+b = config.base**5.2
+c = config.base**runtime.config.exponent
+d = 2**5
+e = 2**~5
+
+# ... but these will be surrounded by whitespace
+f = 2 ** get_exponent()
+g = get_x() ** get_y()
+h = config['base'] ** 2
+```
+
 ### Slices
 
 PEP 8

src/black/linegen.py

@@ -21,8 +21,8 @@
 from black.numerics import normalize_numeric_literal
 from black.strings import get_string_prefix, fix_docstring
 from black.strings import normalize_string_prefix, normalize_string_quotes
-from black.trans import Transformer, CannotTransform, StringMerger
+from black.trans import Transformer, CannotTransform, StringMerger, StringSplitter
-from black.trans import StringSplitter, StringParenWrapper, StringParenStripper
+from black.trans import StringParenWrapper, StringParenStripper, hug_power_op
 from black.mode import Mode, Feature, Preview
 
 from blib2to3.pytree import Node, Leaf
@@ -404,6 +404,9 @@ 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)
 
     for transform in transformers:
         # We are accumulating lines in `result` because we might want to abort

src/black/trans.py

@@ -24,9 +24,9 @@
 import sys
 
 if sys.version_info < (3, 8):
-    from typing_extensions import Final
+    from typing_extensions import Literal, Final
 else:
-    from typing import Final
+    from typing import Literal, Final
 
 from mypy_extensions import trait
 
@@ -71,6 +71,88 @@ def TErr(err_msg: str) -> Err[CannotTransform]:
     return Err(cant_transform)
 
 
+def hug_power_op(line: Line, features: Collection[Feature]) -> Iterator[Line]:
+    """A transformer which normalizes spacing around power operators."""
+
+    # Performance optimization to avoid unnecessary Leaf clones and other ops.
+    for leaf in line.leaves:
+        if leaf.type == token.DOUBLESTAR:
+            break
+    else:
+        raise CannotTransform("No doublestar token was found in the line.")
+
+    def is_simple_lookup(index: int, step: Literal[1, -1]) -> bool:
+        # Brackets and parentheses indicate calls, subscripts, etc. ...
+        # basically stuff that doesn't count as "simple". Only a NAME lookup
+        # or dotted lookup (eg. NAME.NAME) is OK.
+        if step == -1:
+            disallowed = {token.RPAR, token.RSQB}
+        else:
+            disallowed = {token.LPAR, token.LSQB}
+
+        while 0 <= index < len(line.leaves):
+            current = line.leaves[index]
+            if current.type in disallowed:
+                return False
+            if current.type not in {token.NAME, token.DOT} or current.value == "for":
+                # If the current token isn't disallowed, we'll assume this is simple as
+                # only the disallowed tokens are semantically attached to this lookup
+                # expression we're checking. Also, stop early if we hit the 'for' bit
+                # of a comprehension.
+                return True
+
+            index += step
+
+        return True
+
+    def is_simple_operand(index: int, kind: Literal["base", "exponent"]) -> bool:
+        # An operand is considered "simple" if's a NAME, a numeric CONSTANT, a simple
+        # lookup (see above), with or without a preceding unary operator.
+        start = line.leaves[index]
+        if start.type in {token.NAME, token.NUMBER}:
+            return is_simple_lookup(index, step=(1 if kind == "exponent" else -1))
+
+        if start.type in {token.PLUS, token.MINUS, token.TILDE}:
+            if line.leaves[index + 1].type in {token.NAME, token.NUMBER}:
+                # step is always one as bases with a preceding unary op will be checked
+                # for simplicity starting from the next token (so it'll hit the check
+                # above).
+                return is_simple_lookup(index + 1, step=1)
+
+        return False
+
+    leaves: List[Leaf] = []
+    should_hug = False
+    for idx, leaf in enumerate(line.leaves):
+        new_leaf = leaf.clone()
+        if should_hug:
+            new_leaf.prefix = ""
+            should_hug = False
+
+        should_hug = (
+            (0 < idx < len(line.leaves) - 1)
+            and leaf.type == token.DOUBLESTAR
+            and is_simple_operand(idx - 1, kind="base")
+            and line.leaves[idx - 1].value != "lambda"
+            and is_simple_operand(idx + 1, kind="exponent")
+        )
+        if should_hug:
+            new_leaf.prefix = ""
+
+        leaves.append(new_leaf)
+
+    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,
+    )
+
+
 class StringTransformer(ABC):
     """
     An implementation of the Transformer protocol that relies on its

src/black_primer/primer.json

@@ -81,7 +81,7 @@
     },
     "flake8-bugbear": {
       "cli_arguments": ["--experimental-string-processing"],
-      "expect_formatting_changes": false,
+      "expect_formatting_changes": true,
       "git_clone_url": "https://github.com/PyCQA/flake8-bugbear.git",
       "long_checkout": false,
       "py_versions": ["all"]

tests/data/expression.diff

@@ -11,7 +11,17 @@
  True
  False
  1
-@@ -29,63 +29,96 @@
+@@ -21,71 +21,104 @@
+ Name1 or (Name2 and Name3) or Name4
+ Name1 or Name2 and Name3 or Name4
+ v1 << 2
+ 1 >> v2
+ 1 % finished
+-1 + v2 - v3 * 4 ^ 5 ** v6 / 7 // 8
+-((1 + v2) - (v3 * 4)) ^ (((5 ** v6) / 7) // 8)
++1 + v2 - v3 * 4 ^ 5**v6 / 7 // 8
++((1 + v2) - (v3 * 4)) ^ (((5**v6) / 7) // 8)
+ not great
  ~great
  +value
  -1
@@ -88,15 +98,19 @@
 +    *more,
 +]
  {i for i in (1, 2, 3)}
- {(i ** 2) for i in (1, 2, 3)}
+-{(i ** 2) for i in (1, 2, 3)}
 -{(i ** 2) for i, _ in ((1, 'a'), (2, 'b'), (3, 'c'))}
+-{((i ** 2) + j) for i in (1, 2, 3) for j in (1, 2, 3)}
++{(i**2) for i in (1, 2, 3)}
 +{(i**2) for i, _ in ((1, "a"), (2, "b"), (3, "c"))}
- {((i ** 2) + j) for i in (1, 2, 3) for j in (1, 2, 3)}
++{((i**2) + j) for i in (1, 2, 3) for j in (1, 2, 3)}
  [i for i in (1, 2, 3)]
- [(i ** 2) for i in (1, 2, 3)]
+-[(i ** 2) for i in (1, 2, 3)]
 -[(i ** 2) for i, _ in ((1, 'a'), (2, 'b'), (3, 'c'))]
+-[((i ** 2) + j) for i in (1, 2, 3) for j in (1, 2, 3)]
++[(i**2) for i in (1, 2, 3)]
 +[(i**2) for i, _ in ((1, "a"), (2, "b"), (3, "c"))]
- [((i ** 2) + j) for i in (1, 2, 3) for j in (1, 2, 3)]
++[((i**2) + j) for i in (1, 2, 3) for j in (1, 2, 3)]
  {i: 0 for i in (1, 2, 3)}
 -{i: j for i, j in ((1, 'a'), (2, 'b'), (3, 'c'))}
 +{i: j for i, j in ((1, "a"), (2, "b"), (3, "c"))}
@@ -181,10 +195,12 @@
  SomeName
  (Good, Bad, Ugly)
  (i for i in (1, 2, 3))
- ((i ** 2) for i in (1, 2, 3))
+-((i ** 2) for i in (1, 2, 3))
 -((i ** 2) for i, _ in ((1, 'a'), (2, 'b'), (3, 'c')))
+-(((i ** 2) + j) for i in (1, 2, 3) for j in (1, 2, 3))
++((i**2) for i in (1, 2, 3))
 +((i**2) for i, _ in ((1, "a"), (2, "b"), (3, "c")))
- (((i ** 2) + j) for i in (1, 2, 3) for j in (1, 2, 3))
++(((i**2) + j) for i in (1, 2, 3) for j in (1, 2, 3))
  (*starred,)
 -{"id": "1","type": "type","started_at": now(),"ended_at": now() + timedelta(days=10),"priority": 1,"import_session_id": 1,**kwargs}
 +{

tests/data/expression_skip_magic_trailing_comma.diff

@@ -11,7 +11,17 @@
  True
  False
  1
-@@ -29,63 +29,84 @@
+@@ -21,71 +21,92 @@
+ Name1 or (Name2 and Name3) or Name4
+ Name1 or Name2 and Name3 or Name4
+ v1 << 2
+ 1 >> v2
+ 1 % finished
+-1 + v2 - v3 * 4 ^ 5 ** v6 / 7 // 8
+-((1 + v2) - (v3 * 4)) ^ (((5 ** v6) / 7) // 8)
++1 + v2 - v3 * 4 ^ 5**v6 / 7 // 8
++((1 + v2) - (v3 * 4)) ^ (((5**v6) / 7) // 8)
+ not great
  ~great
  +value
  -1
@@ -76,15 +86,19 @@
 +    *more,
 +]
  {i for i in (1, 2, 3)}
- {(i ** 2) for i in (1, 2, 3)}
+-{(i ** 2) for i in (1, 2, 3)}
 -{(i ** 2) for i, _ in ((1, 'a'), (2, 'b'), (3, 'c'))}
+-{((i ** 2) + j) for i in (1, 2, 3) for j in (1, 2, 3)}
++{(i**2) for i in (1, 2, 3)}
 +{(i**2) for i, _ in ((1, "a"), (2, "b"), (3, "c"))}
- {((i ** 2) + j) for i in (1, 2, 3) for j in (1, 2, 3)}
++{((i**2) + j) for i in (1, 2, 3) for j in (1, 2, 3)}
  [i for i in (1, 2, 3)]
- [(i ** 2) for i in (1, 2, 3)]
+-[(i ** 2) for i in (1, 2, 3)]
 -[(i ** 2) for i, _ in ((1, 'a'), (2, 'b'), (3, 'c'))]
+-[((i ** 2) + j) for i in (1, 2, 3) for j in (1, 2, 3)]
++[(i**2) for i in (1, 2, 3)]
 +[(i**2) for i, _ in ((1, "a"), (2, "b"), (3, "c"))]
- [((i ** 2) + j) for i in (1, 2, 3) for j in (1, 2, 3)]
++[((i**2) + j) for i in (1, 2, 3) for j in (1, 2, 3)]
  {i: 0 for i in (1, 2, 3)}
 -{i: j for i, j in ((1, 'a'), (2, 'b'), (3, 'c'))}
 +{i: j for i, j in ((1, "a"), (2, "b"), (3, "c"))}
@@ -164,10 +178,12 @@
  SomeName
  (Good, Bad, Ugly)
  (i for i in (1, 2, 3))
- ((i ** 2) for i in (1, 2, 3))
+-((i ** 2) for i in (1, 2, 3))
 -((i ** 2) for i, _ in ((1, 'a'), (2, 'b'), (3, 'c')))
+-(((i ** 2) + j) for i in (1, 2, 3) for j in (1, 2, 3))
++((i**2) for i in (1, 2, 3))
 +((i**2) for i, _ in ((1, "a"), (2, "b"), (3, "c")))
- (((i ** 2) + j) for i in (1, 2, 3) for j in (1, 2, 3))
++(((i**2) + j) for i in (1, 2, 3) for j in (1, 2, 3))
  (*starred,)
 -{"id": "1","type": "type","started_at": now(),"ended_at": now() + timedelta(days=10),"priority": 1,"import_session_id": 1,**kwargs}
 +{

tests/data/power_op_spacing.py

@@ -0,0 +1,103 @@
+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
+
+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)]
+
+
+# output
+
+
+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
+
+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)]

tests/test_format.py

@@ -48,6 +48,7 @@
     "function2",
     "function_trailing_comma",
     "import_spacing",
+    "power_op_spacing",
     "remove_parens",
     "slices",
     "string_prefixes",