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

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

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 isinstance(bracket_contents, Node) and (
-                if bracket_contents.type != syms.power:
+                bracket_contents.type != syms.power
-                    ensure_visible(opening_bracket)
+                or bracket_contents.children[0].type == token.AWAIT
-                    ensure_visible(closing_bracket)
+                or any(
-                elif (
+                    isinstance(child, Leaf) and child.type == token.DOUBLESTAR
-                    bracket_contents.type == syms.power
+                    for child in bracket_contents.children
-                    and bracket_contents.children[0].type == token.AWAIT
+                )
-                ):
+            ):
-                    ensure_visible(opening_bracket)
+                ensure_visible(opening_bracket)
-                    ensure_visible(closing_bracket)
+                ensure_visible(closing_bracket)
-                    # If we are in a nested await then recurse down.
-                    remove_await_parens(bracket_contents)
 
 
 def _maybe_wrap_cms_in_parens(

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]