Use recursion to remove redundant brackets from nested awaits

src/black/linegen.py

@@ -3,7 +3,7 @@
 """
 from functools import partial, wraps
 import sys
-from typing import Collection, Iterator, List, Optional, Set, Union
+from typing import Collection, Iterator, List, Optional, Set, Union, cast
 
 from black.nodes import WHITESPACE, RARROW, STATEMENT, STANDALONE_COMMENT
 from black.nodes import ASSIGNMENTS, OPENING_BRACKETS, CLOSING_BRACKETS
@@ -226,42 +226,8 @@ def visit_power(self, node: Node) -> Iterator[Line]:
             ):
                 wrap_in_parentheses(node, leaf)
 
-        if (
-            Preview.remove_redundant_parens in self.mode
-            and node.children[0].type == token.AWAIT
-            and len(node.children) > 1
-        ):
-            if (
-                node.children[1].type == syms.atom
-                and node.children[1].children[0].type == token.LPAR
-            ):
-                if maybe_make_parens_invisible_in_atom(
-                    node.children[1],
-                    parent=node,
-                    remove_brackets_around_comma=True,
-                ):
-                    wrap_in_parentheses(node, node.children[1], visible=False)
-                if (
-                    isinstance(node.children[1].children[0], Leaf)
-                    and isinstance(node.children[1].children[-1], Leaf)
-                    and (
-                        node.children[1].children[1].type != syms.power
-                        or (
-                            node.children[1].children[1].type == syms.power
-                            and node.children[1].children[1].children[0].type
-                            == token.AWAIT
-                        )
-                    )
-                ):
-                    # Since await is an expression we shouldn't remove
-                    # brackets in cases where this would change
-                    # the AST due to operator precedence.
-                    # Therefore we only aim to remove brackets around
-                    # power nodes that aren't also await expressions themselves.
-                    # https://peps.python.org/pep-0492/#updated-operator-precedence-table
-                    # N.B. We've still removed any redundant nested brackets though :)
-                    ensure_visible(node.children[1].children[0])
-                    ensure_visible(node.children[1].children[-1])
+        if Preview.remove_redundant_parens in self.mode:
+            remove_await_parens(node)
 
         yield from self.visit_default(node)
 
@@ -931,6 +897,42 @@ def normalize_invisible_parens(
         )
 
 
+def remove_await_parens(node: Node) -> None:
+    if node.children[0].type == token.AWAIT and len(node.children) > 1:
+        if (
+            node.children[1].type == syms.atom
+            and node.children[1].children[0].type == token.LPAR
+        ):
+            if maybe_make_parens_invisible_in_atom(
+                node.children[1],
+                parent=node,
+                remove_brackets_around_comma=True,
+            ):
+                wrap_in_parentheses(node, node.children[1], visible=False)
+
+            # Since await is an expression we shouldn't remove
+            # brackets in cases where this would change
+            # the AST due to operator precedence.
+            # Therefore we only aim to remove brackets around
+            # power nodes that aren't also await expressions themselves.
+            # https://peps.python.org/pep-0492/#updated-operator-precedence-table
+            # N.B. We've still removed any redundant nested brackets though :)
+            opening_bracket = cast(Leaf, node.children[1].children[0])
+            closing_bracket = cast(Leaf, node.children[1].children[-1])
+            bracket_contents = cast(Node, node.children[1].children[1])
+            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)
+
+
 def remove_with_parens(node: Node, parent: Node) -> None:
     """Recursively hide optional parens in `with` statements."""
     # Removing all unnecessary parentheses in with statements in one pass is a tad

tests/data/remove_await_parens.py

@@ -64,6 +64,13 @@ async def main():
 async def main():
     await (await asyncio.sleep(1))
 
+# It's brackets all the way down...
+async def main():
+    await (await (asyncio.sleep(1)))
+
+async def main():
+    await (await (await (await (await (asyncio.sleep(1))))))
+
 # output
 import asyncio
 
@@ -136,3 +143,12 @@ async def main():
 
 async def main():
     await (await asyncio.sleep(1))
+
+
+# It's brackets all the way down...
+async def main():
+    await (await asyncio.sleep(1))
+
+
+async def main():
+    await (await (await (await (await asyncio.sleep(1)))))