+
+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 = 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)
+
+
+def _maybe_wrap_cms_in_parens(
+ node: Node, mode: Mode, features: Collection[Feature]
+) -> None:
+ """When enabled and safe, wrap the multiple context managers in invisible parens.
+
+ It is only safe when `features` contain Feature.PARENTHESIZED_CONTEXT_MANAGERS.
+ """
+ if (
+ Feature.PARENTHESIZED_CONTEXT_MANAGERS not in features
+ or Preview.wrap_multiple_context_managers_in_parens not in mode
+ or len(node.children) <= 2
+ # If it's an atom, it's already wrapped in parens.
+ or node.children[1].type == syms.atom
+ ):
+ return
+ colon_index: Optional[int] = None
+ for i in range(2, len(node.children)):
+ if node.children[i].type == token.COLON:
+ colon_index = i
+ break
+ if colon_index is not None:
+ lpar = Leaf(token.LPAR, "")
+ rpar = Leaf(token.RPAR, "")
+ context_managers = node.children[1:colon_index]
+ for child in context_managers:
+ child.remove()
+ # After wrapping, the with_stmt will look like this:
+ # with_stmt
+ # NAME 'with'
+ # atom
+ # LPAR ''
+ # testlist_gexp
+ # ... <-- context_managers
+ # /testlist_gexp
+ # RPAR ''
+ # /atom
+ # COLON ':'
+ new_child = Node(
+ syms.atom, [lpar, Node(syms.testlist_gexp, context_managers), rpar]
+ )
+ node.insert_child(1, new_child)
+
+
+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
+ # complex as different variations of bracketed statements result in pretty
+ # different parse trees:
+ #
+ # with (open("file")) as f: # this is an asexpr_test
+ # ...
+ #
+ # with (open("file") as f): # this is an atom containing an
+ # ... # asexpr_test
+ #
+ # with (open("file")) as f, (open("file")) as f: # this is asexpr_test, COMMA,
+ # ... # asexpr_test
+ #
+ # with (open("file") as f, open("file") as f): # an atom containing a
+ # ... # testlist_gexp which then
+ # # contains multiple asexpr_test(s)
+ if node.type == syms.atom:
+ if maybe_make_parens_invisible_in_atom(
+ node,
+ parent=parent,
+ remove_brackets_around_comma=True,
+ ):
+ wrap_in_parentheses(parent, node, visible=False)
+ if isinstance(node.children[1], Node):
+ remove_with_parens(node.children[1], node)
+ elif node.type == syms.testlist_gexp:
+ for child in node.children:
+ if isinstance(child, Node):
+ remove_with_parens(child, node)
+ elif node.type == syms.asexpr_test and not any(
+ leaf.type == token.COLONEQUAL for leaf in node.leaves()
+ ):
+ if maybe_make_parens_invisible_in_atom(
+ node.children[0],
+ parent=node,
+ remove_brackets_around_comma=True,
+ ):
+ wrap_in_parentheses(node, node.children[0], visible=False)
+
+
+def maybe_make_parens_invisible_in_atom(
+ node: LN,
+ parent: LN,
+ remove_brackets_around_comma: bool = False,
+) -> bool: