From c54c213d6a3132986feede0cf0525f5bae5b43d6 Mon Sep 17 00:00:00 2001
From: Jelle Zijlstra <jelle.zijlstra@gmail.com>
Date: Thu, 2 Nov 2023 20:42:11 -0700
Subject: [PATCH] Fix crash on await (a ** b) (#3994)

---
 CHANGES.md                              |  2 ++
 src/black/linegen.py                    | 22 ++++++++++------------
 tests/data/cases/remove_await_parens.py | 19 +++++++++++++++++++
 3 files changed, 31 insertions(+), 12 deletions(-)

diff --git a/CHANGES.md b/CHANGES.md
index f365f1c239b..5ce37943693 100644
--- 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
index 121c6e314fe..b13b95d9b31 100644
--- 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
index 8c7223d2f39..073150c5f08 100644
--- 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]