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

src/black/__init__.py

@@ -272,6 +272,7 @@ class Feature(Enum):
 class Mode:
     target_versions: Set[TargetVersion] = field(default_factory=set)
     line_length: int = DEFAULT_LINE_LENGTH
+    doublestar_math_op: bool = False
     string_normalization: bool = True
     magic_trailing_comma: bool = True
     experimental_string_processing: bool = False
@@ -432,6 +433,12 @@ def validate_regex(
     is_flag=True,
     help="Don't use trailing commas as a reason to split lines.",
 )
+@click.option(
+    "-DS",
+    "--skip-double-star-op-spaces",
+    is_flag=True,
+    help="Don't surround ** (doublestar) with spaces if math op.",
+)
 @click.option(
     "--experimental-string-processing",
     is_flag=True,
@@ -571,6 +578,7 @@ def main(
     fast: bool,
     pyi: bool,
     skip_string_normalization: bool,
+    skip_double_star_op_spaces: bool,
     skip_magic_trailing_comma: bool,
     experimental_string_processing: bool,
     quiet: bool,
@@ -596,6 +604,7 @@ def main(
         is_pyi=pyi,
         string_normalization=not skip_string_normalization,
         magic_trailing_comma=not skip_magic_trailing_comma,
+        doublestar_math_op=skip_double_star_op_spaces,
         experimental_string_processing=experimental_string_processing,
     )
     if config and verbose:
@@ -1078,6 +1087,8 @@ def f(
         for line in transform_line(
             current_line, mode=mode, features=split_line_features
         ):
+            if line.mode.doublestar_math_op:
+                line = fix_doublestar_in_op(line)
             dst_contents.append(str(line))
     return "".join(dst_contents)
 
@@ -1844,6 +1855,34 @@ 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] = []
+    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 = ""
+        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 = ""
+        leaves.append(new_leaf)
+
+    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,15 @@
+#!/usr/bin/env python3.7
+
+
+def function(**kwargs):
+    t = a**2 + b**3
+
+
+# output
+
+
+#!/usr/bin/env python3.7
+
+
+def function(**kwargs):
+    t = a**2 + b**3

tests/test_black.py

@@ -407,6 +407,17 @@ 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_two_magic_commas(self) -> None:
+        source, expected = read_data("doublestar_no_spaces")
+        mode = replace(
+            DEFAULT_MODE, target_versions=PY36_VERSIONS, doublestar_math_op=True
+        )
+        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")