Wrap multiple context managers in parentheses when targeting Python 3.9+

CHANGES.md

@@ -35,6 +35,7 @@
 - Fix two crashes in preview style involving edge cases with docstrings (#3451)
 - Exclude string type annotations from improved string processing; fix crash when the
   return type annotation is stringified and spans across multiple lines (#3462)
+- Wrap multiple context managers in parentheses when targeting Python 3.9+ (#3489)
 - Fix several crashes in preview style with walrus operators used in `with` statements
   or tuples (#3473)
 

src/black/__init__.py

@@ -1096,8 +1096,13 @@ def _format_str_once(src_contents: str, *, mode: Mode) -> str:
         future_imports = get_future_imports(src_node)
         versions = detect_target_versions(src_node, future_imports=future_imports)
 
+    context_manager_features = {
+        feature
+        for feature in {Feature.PARENTHESIZED_CONTEXT_MANAGERS}
+        if supports_feature(versions, feature)
+    }
     normalize_fmt_off(src_node, preview=mode.preview)
-    lines = LineGenerator(mode=mode)
+    lines = LineGenerator(mode=mode, features=context_manager_features)
     elt = EmptyLineTracker(mode=mode)
     split_line_features = {
         feature
@@ -1159,6 +1164,10 @@ def get_features_used(  # noqa: C901
     - relaxed decorator syntax;
     - usage of __future__ flags (annotations);
     - print / exec statements;
+    - parenthesized context managers;
+    - match statements;
+    - except* clause;
+    - variadic generics;
     """
     features: Set[Feature] = set()
     if future_imports:
@@ -1234,6 +1243,23 @@ def get_features_used(  # noqa: C901
         ):
             features.add(Feature.ANN_ASSIGN_EXTENDED_RHS)
 
+        elif (
+            n.type == syms.with_stmt
+            and len(n.children) > 2
+            and n.children[1].type == syms.atom
+        ):
+            atom_children = n.children[1].children
+            if (
+                len(atom_children) == 3
+                and atom_children[0].type == token.LPAR
+                and atom_children[1].type == syms.testlist_gexp
+                and atom_children[2].type == token.RPAR
+            ):
+                features.add(Feature.PARENTHESIZED_CONTEXT_MANAGERS)
+
+        elif n.type == syms.match_stmt:
+            features.add(Feature.PATTERN_MATCHING)
+
         elif (
             n.type == syms.except_clause
             and len(n.children) >= 2

src/black/linegen.py

@@ -90,8 +90,9 @@ class LineGenerator(Visitor[Line]):
     in ways that will no longer stringify to valid Python code on the tree.
     """
 
-    def __init__(self, mode: Mode) -> None:
+    def __init__(self, mode: Mode, features: Collection[Feature]) -> None:
         self.mode = mode
+        self.features = features
         self.current_line: Line
         self.__post_init__()
 
@@ -191,7 +192,9 @@ def visit_stmt(
         `parens` holds a set of string leaf values immediately after which
         invisible parens should be put.
         """
-        normalize_invisible_parens(node, parens_after=parens, preview=self.mode.preview)
+        normalize_invisible_parens(
+            node, parens_after=parens, mode=self.mode, features=self.features
+        )
         for child in node.children:
             if is_name_token(child) and child.value in keywords:
                 yield from self.line()
@@ -244,7 +247,9 @@ def visit_funcdef(self, node: Node) -> Iterator[Line]:
 
     def visit_match_case(self, node: Node) -> Iterator[Line]:
         """Visit either a match or case statement."""
-        normalize_invisible_parens(node, parens_after=set(), preview=self.mode.preview)
+        normalize_invisible_parens(
+            node, parens_after=set(), mode=self.mode, features=self.features
+        )
 
         yield from self.line()
         for child in node.children:
@@ -1090,7 +1095,7 @@ def normalize_prefix(leaf: Leaf, *, inside_brackets: bool) -> None:
 
 
 def normalize_invisible_parens(
-    node: Node, parens_after: Set[str], *, preview: bool
+    node: Node, parens_after: Set[str], *, mode: Mode, features: Collection[Feature]
 ) -> None:
     """Make existing optional parentheses invisible or create new ones.
 
@@ -1100,17 +1105,24 @@ def normalize_invisible_parens(
     Standardizes on visible parentheses for single-element tuples, and keeps
     existing visible parentheses for other tuples and generator expressions.
     """
-    for pc in list_comments(node.prefix, is_endmarker=False, preview=preview):
+    for pc in list_comments(node.prefix, is_endmarker=False, preview=mode.preview):
         if pc.value in FMT_OFF:
             # This `node` has a prefix with `# fmt: off`, don't mess with parens.
             return
+
+    # The multiple context managers grammar has a different pattern, thus this is
+    # separate from the for-loop below. This possibly wraps them in invisible parens,
+    # and later will be removed in remove_with_parens when needed.
+    if node.type == syms.with_stmt:
+        _maybe_wrap_cms_in_parens(node, mode, features)
+
     check_lpar = False
     for index, child in enumerate(list(node.children)):
         # Fixes a bug where invisible parens are not properly stripped from
         # assignment statements that contain type annotations.
         if isinstance(child, Node) and child.type == syms.annassign:
             normalize_invisible_parens(
-                child, parens_after=parens_after, preview=preview
+                child, parens_after=parens_after, mode=mode, features=features
             )
 
         # Add parentheses around long tuple unpacking in assignments.
@@ -1123,7 +1135,7 @@ def normalize_invisible_parens(
 
         if check_lpar:
             if (
-                preview
+                mode.preview
                 and child.type == syms.atom
                 and node.type == syms.for_stmt
                 and isinstance(child.prev_sibling, Leaf)
@@ -1136,7 +1148,9 @@ def normalize_invisible_parens(
                     remove_brackets_around_comma=True,
                 ):
                     wrap_in_parentheses(node, child, visible=False)
-            elif preview and isinstance(child, Node) and node.type == syms.with_stmt:
+            elif (
+                mode.preview and isinstance(child, Node) and node.type == syms.with_stmt
+            ):
                 remove_with_parens(child, node)
             elif child.type == syms.atom:
                 if maybe_make_parens_invisible_in_atom(
@@ -1147,17 +1161,7 @@ def normalize_invisible_parens(
             elif is_one_tuple(child):
                 wrap_in_parentheses(node, child, visible=True)
             elif node.type == syms.import_from:
-                # "import from" nodes store parentheses directly as part of
-                # the statement
-                if is_lpar_token(child):
-                    assert is_rpar_token(node.children[-1])
-                    # make parentheses invisible
-                    child.value = ""
-                    node.children[-1].value = ""
-                elif child.type != token.STAR:
-                    # insert invisible parentheses
-                    node.insert_child(index, Leaf(token.LPAR, ""))
-                    node.append_child(Leaf(token.RPAR, ""))
+                _normalize_import_from(node, child, index)
                 break
             elif (
                 index == 1
@@ -1172,13 +1176,27 @@ def normalize_invisible_parens(
             elif not (isinstance(child, Leaf) and is_multiline_string(child)):
                 wrap_in_parentheses(node, child, visible=False)
 
-        comma_check = child.type == token.COMMA if preview else False
+        comma_check = child.type == token.COMMA if mode.preview else False
 
         check_lpar = isinstance(child, Leaf) and (
             child.value in parens_after or comma_check
         )
 
 
+def _normalize_import_from(parent: Node, child: LN, index: int) -> None:
+    # "import from" nodes store parentheses directly as part of
+    # the statement
+    if is_lpar_token(child):
+        assert is_rpar_token(parent.children[-1])
+        # make parentheses invisible
+        child.value = ""
+        parent.children[-1].value = ""
+    elif child.type != token.STAR:
+        # insert invisible parentheses
+        parent.insert_child(index, Leaf(token.LPAR, ""))
+        parent.append_child(Leaf(token.RPAR, ""))
+
+
 def remove_await_parens(node: Node) -> None:
     if node.children[0].type == token.AWAIT and len(node.children) > 1:
         if (
@@ -1215,6 +1233,49 @@ def remove_await_parens(node: Node) -> None:
                 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

src/black/mode.py

@@ -50,6 +50,7 @@ class Feature(Enum):
     EXCEPT_STAR = 14
     VARIADIC_GENERICS = 15
     DEBUG_F_STRINGS = 16
+    PARENTHESIZED_CONTEXT_MANAGERS = 17
     FORCE_OPTIONAL_PARENTHESES = 50
 
     # __future__ flags
@@ -106,6 +107,7 @@ class Feature(Enum):
         Feature.POS_ONLY_ARGUMENTS,
         Feature.UNPACKING_ON_FLOW,
         Feature.ANN_ASSIGN_EXTENDED_RHS,
+        Feature.PARENTHESIZED_CONTEXT_MANAGERS,
     },
     TargetVersion.PY310: {
         Feature.F_STRINGS,
@@ -120,6 +122,7 @@ class Feature(Enum):
         Feature.POS_ONLY_ARGUMENTS,
         Feature.UNPACKING_ON_FLOW,
         Feature.ANN_ASSIGN_EXTENDED_RHS,
+        Feature.PARENTHESIZED_CONTEXT_MANAGERS,
         Feature.PATTERN_MATCHING,
     },
     TargetVersion.PY311: {
@@ -135,6 +138,7 @@ class Feature(Enum):
         Feature.POS_ONLY_ARGUMENTS,
         Feature.UNPACKING_ON_FLOW,
         Feature.ANN_ASSIGN_EXTENDED_RHS,
+        Feature.PARENTHESIZED_CONTEXT_MANAGERS,
         Feature.PATTERN_MATCHING,
         Feature.EXCEPT_STAR,
         Feature.VARIADIC_GENERICS,
@@ -164,6 +168,7 @@ class Preview(Enum):
     parenthesize_conditional_expressions = auto()
     skip_magic_trailing_comma_in_subscript = auto()
     wrap_long_dict_values_in_parens = auto()
+    wrap_multiple_context_managers_in_parens = auto()
 
 
 class Deprecated(UserWarning):

tests/data/preview_context_managers/auto_detect/features_3_10.py

@@ -0,0 +1,35 @@
+# This file uses pattern matching introduced in Python 3.10.
+
+
+match http_code:
+    case 404:
+        print("Not found")
+
+
+with \
+     make_context_manager1() as cm1, \
+     make_context_manager2() as cm2, \
+     make_context_manager3() as cm3, \
+     make_context_manager4() as cm4 \
+:
+    pass
+
+
+# output
+
+
+# This file uses pattern matching introduced in Python 3.10.
+
+
+match http_code:
+    case 404:
+        print("Not found")
+
+
+with (
+    make_context_manager1() as cm1,
+    make_context_manager2() as cm2,
+    make_context_manager3() as cm3,
+    make_context_manager4() as cm4,
+):
+    pass

tests/data/preview_context_managers/auto_detect/features_3_11.py

@@ -0,0 +1,37 @@
+# This file uses except* clause in Python 3.11.
+
+
+try:
+    some_call()
+except* Error as e:
+    pass
+
+
+with \
+     make_context_manager1() as cm1, \
+     make_context_manager2() as cm2, \
+     make_context_manager3() as cm3, \
+     make_context_manager4() as cm4 \
+:
+    pass
+
+
+# output
+
+
+# This file uses except* clause in Python 3.11.
+
+
+try:
+    some_call()
+except* Error as e:
+    pass
+
+
+with (
+    make_context_manager1() as cm1,
+    make_context_manager2() as cm2,
+    make_context_manager3() as cm3,
+    make_context_manager4() as cm4,
+):
+    pass

tests/data/preview_context_managers/auto_detect/features_3_8.py

@@ -0,0 +1,30 @@
+# This file doesn't use any Python 3.9+ only grammars.
+
+
+# Make sure parens around a single context manager don't get autodetected as
+# Python 3.9+.
+with (a):
+    pass
+
+
+with \
+     make_context_manager1() as cm1, \
+     make_context_manager2() as cm2, \
+     make_context_manager3() as cm3, \
+     make_context_manager4() as cm4 \
+:
+    pass
+
+
+# output
+# This file doesn't use any Python 3.9+ only grammars.
+
+
+# Make sure parens around a single context manager don't get autodetected as
+# Python 3.9+.
+with a:
+    pass
+
+
+with make_context_manager1() as cm1, make_context_manager2() as cm2, make_context_manager3() as cm3, make_context_manager4() as cm4:
+    pass