Fix crash on await (a ** b) (#3994)

psf · Nov 3, 2023 · c54c213 · c54c213
1 parent e2f2bd0
commit c54c213
Show file tree

Hide file tree

Showing 3 changed files with 31 additions and 12 deletions.
diff --git a/CHANGES.md b/CHANGES.md
@@ -11,6 +11,8 @@
 - Fix crash on formatting bytes strings that look like docstrings (#4003)
 - Fix crash when whitespace followed a backslash before newline in a docstring (#4008)
 
+- Fix crash on formatting code like `await (a ** b)` (#3994)
+
 ### Preview style
 
 - Multiline dictionaries and lists that are the sole argument to a function are now

diff --git a/src/black/linegen.py b/src/black/linegen.py
@@ -1352,18 +1352,16 @@ def remove_await_parens(node: Node) -> None:
             opening_bracket = cast(Leaf, node.children[1].children[0])
             closing_bracket = cast(Leaf, node.children[1].children[-1])
             bracket_contents = node.children[1].children[1]
-            if isinstance(bracket_contents, Node):
-                if bracket_contents.type != syms.power:
-                    ensure_visible(opening_bracket)
-                    ensure_visible(closing_bracket)
-                elif (
-                    bracket_contents.type == syms.power
-                    and bracket_contents.children[0].type == token.AWAIT
-                ):
-                    ensure_visible(opening_bracket)
-                    ensure_visible(closing_bracket)
-                    # If we are in a nested await then recurse down.
-                    remove_await_parens(bracket_contents)
+            if isinstance(bracket_contents, Node) and (
+                bracket_contents.type != syms.power
+                or bracket_contents.children[0].type == token.AWAIT
+                or any(
+                    isinstance(child, Leaf) and child.type == token.DOUBLESTAR
+                    for child in bracket_contents.children
+                )
+            ):
+                ensure_visible(opening_bracket)
+                ensure_visible(closing_bracket)
 
 
 def _maybe_wrap_cms_in_parens(

diff --git a/tests/data/cases/remove_await_parens.py b/tests/data/cases/remove_await_parens.py
@@ -80,6 +80,15 @@ async def main():
 async def main():
     await (yield)
 
+async def main():
+    await (a ** b)
+    await (a[b] ** c)
+    await (a ** b[c])
+    await ((a + b) ** (c + d))
+    await (a + b)
+    await (a[b])
+    await (a[b ** c])
+
 # output
 import asyncio
 
@@ -174,3 +183,13 @@ async def main():
 
 async def main():
     await (yield)
+
+
+async def main():
+    await (a**b)
+    await (a[b] ** c)
+    await (a ** b[c])
+    await ((a + b) ** (c + d))
+    await (a + b)
+    await a[b]
+    await a[b**c]