Allow no space in doublestar when math operation (#538) (#2095)

CHANGES.md

@@ -4,6 +4,8 @@
 
 #### _Black_
 
+- `Black` wraps `**` if in math expression with more than 1 op (#2095)
+
 - `Black` now respects `--skip-string-normalization` when normalizing multiline
   docstring quotes (#1637)
 

src/black/__init__.py

@@ -1078,6 +1078,7 @@ def f(
         for line in transform_line(
             current_line, mode=mode, features=split_line_features
         ):
+            line = fix_doublestar_in_op(line)
             dst_contents.append(str(line))
     return "".join(dst_contents)
 
@@ -1844,6 +1845,40 @@ def __bool__(self) -> bool:
         return bool(self.leaves or self.comments)
 
 
+def fix_doublestar_in_op(line: Line) -> Line:
+    """Clone line without spaces on doublestar op if is math op."""
+    leaves: List[Leaf] = []
+    candidate_rep_count: int = 0
+    for idx, leaf in enumerate(line.leaves):
+        new_leaf = leaf.clone()
+        if new_leaf.type == token.DOUBLESTAR:
+            if idx > 0 and line.leaves[idx - 1].type in {token.NAME, token.NUMBER}:
+                new_leaf.prefix = ""
+                candidate_rep_count += 1
+        if (
+            idx > 1
+            and line.leaves[idx - 1].type == token.DOUBLESTAR
+            and line.leaves[idx - 2].type in {token.NAME, token.NUMBER}
+        ):
+            new_leaf.prefix = ""
+            candidate_rep_count += 1
+        leaves.append(new_leaf)
+
+    if candidate_rep_count <= 2:
+        leaves = line.leaves
+
+    return 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,
+    )
+
+
 @dataclass
 class EmptyLineTracker:
     """Provides a stateful method that returns the number of potential extra

tests/data/doublestar_no_spaces.py

@@ -0,0 +1,43 @@
+#!/usr/bin/env python3.7
+
+
+def function(**kwargs):
+    t = a**2 + b**3
+
+
+def function_no_spaces():
+    return t**2
+
+
+def function_replace_spaces(**kwargs):
+    t = a **2 + b** 3 + c ** 4
+
+
+def function_dont_replace_spaces():
+    t = t ** 2
+    {**a, **b, **c}
+
+
+
+# output
+
+
+#!/usr/bin/env python3.7
+
+
+def function(**kwargs):
+    t = a**2 + b**3
+
+
+def function_no_spaces():
+    return t ** 2
+
+
+def function_replace_spaces(**kwargs):
+    t = a**2 + b**3 + c**4
+
+
+def function_dont_replace_spaces():
+    t = t ** 2
+    {**a, **b, **c}
+

tests/data/expression.diff

@@ -19,7 +19,7 @@
  (~int) and (not ((v1 ^ (123 + v2)) | True))
 -+really ** -confusing ** ~operator ** -precedence
 -flags & ~ select.EPOLLIN and waiters.write_task is not None
-++(really ** -(confusing ** ~(operator ** -precedence)))
+++(really**-(confusing**~(operator**-precedence)))
 +flags & ~select.EPOLLIN and waiters.write_task is not None
  lambda arg: None
  lambda a=True: a

tests/data/expression.py

@@ -290,7 +290,7 @@ async def f():
 -1
 ~int and not v1 ^ 123 + v2 | True
 (~int) and (not ((v1 ^ (123 + v2)) | True))
-+(really ** -(confusing ** ~(operator ** -precedence)))
++(really**-(confusing**~(operator**-precedence)))
 flags & ~select.EPOLLIN and waiters.write_task is not None
 lambda arg: None
 lambda a=True: a

tests/data/expression_skip_magic_trailing_comma.diff

@@ -19,7 +19,7 @@
  (~int) and (not ((v1 ^ (123 + v2)) | True))
 -+really ** -confusing ** ~operator ** -precedence
 -flags & ~ select.EPOLLIN and waiters.write_task is not None
-++(really ** -(confusing ** ~(operator ** -precedence)))
+++(really**-(confusing**~(operator**-precedence)))
 +flags & ~select.EPOLLIN and waiters.write_task is not None
  lambda arg: None
  lambda a=True: a

tests/data/pep_572.py

@@ -4,7 +4,7 @@
     pass
 if match := pattern.search(data):
     pass
-[y := f(x), y ** 2, y ** 3]
+[y := f(x), y**2, y**3]
 filtered_data = [y for x in data if (y := f(x)) is None]
 (y := f(x))
 y0 = (y1 := f(x))

tests/test_black.py

@@ -407,6 +407,15 @@ def test_numeric_literals_ignoring_underscores(self) -> None:
         black.assert_equivalent(source, actual)
         black.assert_stable(source, actual, mode)
 
+    @patch("black.dump_to_file", dump_to_stderr)
+    def test_doublestar_math_op_no_spaces(self) -> None:
+        source, expected = read_data("doublestar_no_spaces")
+        mode = replace(DEFAULT_MODE, target_versions=PY36_VERSIONS)
+        actual = fs(source, mode=mode)
+        self.assertFormatEqual(expected, actual)
+        black.assert_equivalent(source, actual)
+        black.assert_stable(source, actual, mode)
+
     def test_skip_magic_trailing_comma(self) -> None:
         source, _ = read_data("expression.py")
         expected, _ = read_data("expression_skip_magic_trailing_comma.diff")