Consistently wrap two context managers in parens (in --preview). (#3589)

Co-authored-by: Jelle Zijlstra <jelle.zijlstra@gmail.com>

CHANGES.md

@@ -16,6 +16,8 @@
 
 - Add trailing commas to collection literals even if there's a comment after the last
   entry (#3393)
+- `with` statements that contain two context managers will be consistently wrapped in
+  parentheses (#3589)
 
 ### Configuration
 

src/black/linegen.py

@@ -2,7 +2,7 @@
 Generating lines of code.
 """
 import sys
-from dataclasses import dataclass, replace
+from dataclasses import replace
 from enum import Enum, auto
 from functools import partial, wraps
 from typing import Collection, Iterator, List, Optional, Set, Union, cast
@@ -16,6 +16,7 @@
 from black.comments import FMT_OFF, generate_comments, list_comments
 from black.lines import (
     Line,
+    RHSResult,
     append_leaves,
     can_be_split,
     can_omit_invisible_parens,
@@ -647,17 +648,6 @@ def left_hand_split(
             yield result
 
 
-@dataclass
-class _RHSResult:
-    """Intermediate split result from a right hand split."""
-
-    head: Line
-    body: Line
-    tail: Line
-    opening_bracket: Leaf
-    closing_bracket: Leaf
-
-
 def right_hand_split(
     line: Line,
     mode: Mode,
@@ -681,7 +671,7 @@ def right_hand_split(
 def _first_right_hand_split(
     line: Line,
     omit: Collection[LeafID] = (),
-) -> _RHSResult:
+) -> RHSResult:
     """Split the line into head, body, tail starting with the last bracket pair.
 
     Note: this function should not have side effects. It's relied upon by
@@ -723,11 +713,11 @@ def _first_right_hand_split(
         tail_leaves, line, opening_bracket, component=_BracketSplitComponent.tail
     )
     bracket_split_succeeded_or_raise(head, body, tail)
-    return _RHSResult(head, body, tail, opening_bracket, closing_bracket)
+    return RHSResult(head, body, tail, opening_bracket, closing_bracket)
 
 
 def _maybe_split_omitting_optional_parens(
-    rhs: _RHSResult,
+    rhs: RHSResult,
     line: Line,
     mode: Mode,
     features: Collection[Feature] = (),
@@ -747,11 +737,11 @@ def _maybe_split_omitting_optional_parens(
         # there are no standalone comments in the body
         and not rhs.body.contains_standalone_comments(0)
         # and we can actually remove the parens
-        and can_omit_invisible_parens(rhs.body, mode.line_length)
+        and can_omit_invisible_parens(rhs, mode.line_length)
     ):
         omit = {id(rhs.closing_bracket), *omit}
         try:
-            # The _RHSResult Omitting Optional Parens.
+            # The RHSResult Omitting Optional Parens.
             rhs_oop = _first_right_hand_split(line, omit=omit)
             if not (
                 Preview.prefer_splitting_right_hand_side_of_assignments in line.mode
@@ -803,7 +793,7 @@ def _maybe_split_omitting_optional_parens(
             yield result
 
 
-def _prefer_split_rhs_oop(rhs_oop: _RHSResult, mode: Mode) -> bool:
+def _prefer_split_rhs_oop(rhs_oop: RHSResult, mode: Mode) -> bool:
     """
     Returns whether we should prefer the result from a split omitting optional parens.
     """

src/black/lines.py

@@ -15,7 +15,7 @@
     cast,
 )
 
-from black.brackets import DOT_PRIORITY, BracketTracker
+from black.brackets import COMMA_PRIORITY, DOT_PRIORITY, BracketTracker
 from black.mode import Mode, Preview
 from black.nodes import (
     BRACKETS,
@@ -28,6 +28,7 @@
     is_multiline_string,
     is_one_sequence_between,
     is_type_comment,
+    is_with_stmt,
     replace_child,
     syms,
     whitespace,
@@ -122,6 +123,11 @@ def is_import(self) -> bool:
         """Is this an import line?"""
         return bool(self) and is_import(self.leaves[0])
 
+    @property
+    def is_with_stmt(self) -> bool:
+        """Is this a with_stmt line?"""
+        return bool(self) and is_with_stmt(self.leaves[0])
+
     @property
     def is_class(self) -> bool:
         """Is this line a class definition?"""
@@ -449,6 +455,17 @@ def __bool__(self) -> bool:
         return bool(self.leaves or self.comments)
 
 
+@dataclass
+class RHSResult:
+    """Intermediate split result from a right hand split."""
+
+    head: Line
+    body: Line
+    tail: Line
+    opening_bracket: Leaf
+    closing_bracket: Leaf
+
+
 @dataclass
 class LinesBlock:
     """Class that holds information about a block of formatted lines.
@@ -830,25 +847,42 @@ def can_be_split(line: Line) -> bool:
 
 
 def can_omit_invisible_parens(
-    line: Line,
+    rhs: RHSResult,
     line_length: int,
 ) -> bool:
-    """Does `line` have a shape safe to reformat without optional parens around it?
+    """Does `rhs.body` have a shape safe to reformat without optional parens around it?
 
     Returns True for only a subset of potentially nice looking formattings but
     the point is to not return false positives that end up producing lines that
     are too long.
     """
+    line = rhs.body
     bt = line.bracket_tracker
     if not bt.delimiters:
         # Without delimiters the optional parentheses are useless.
         return True
 
     max_priority = bt.max_delimiter_priority()
-    if bt.delimiter_count_with_priority(max_priority) > 1:
+    delimiter_count = bt.delimiter_count_with_priority(max_priority)
+    if delimiter_count > 1:
         # With more than one delimiter of a kind the optional parentheses read better.
         return False
 
+    if delimiter_count == 1:
+        if (
+            Preview.wrap_multiple_context_managers_in_parens in line.mode
+            and max_priority == COMMA_PRIORITY
+            and rhs.head.is_with_stmt
+        ):
+            # For two context manager with statements, the optional parentheses read
+            # better. In this case, `rhs.body` is the context managers part of
+            # the with statement. `rhs.head` is the `with (` part on the previous
+            # line.
+            return False
+        # Otherwise it may also read better, but we don't do it today and requires
+        # careful considerations for all possible cases. See
+        # https://github.com/psf/black/issues/2156.
+
     if max_priority == DOT_PRIORITY:
         # A single stranded method call doesn't require optional parentheses.
         return True

src/black/nodes.py

@@ -789,6 +789,16 @@ def is_import(leaf: Leaf) -> bool:
     )
 
 
+def is_with_stmt(leaf: Leaf) -> bool:
+    """Return True if the given leaf starts a with statement."""
+    return bool(
+        leaf.type == token.NAME
+        and leaf.value == "with"
+        and leaf.parent
+        and leaf.parent.type == syms.with_stmt
+    )
+
+
 def is_type_comment(leaf: Leaf, suffix: str = "") -> bool:
     """Return True if the given leaf is a special comment.
     Only returns true for type comments for now."""

tests/data/preview_context_managers/auto_detect/features_3_8.py

@@ -16,6 +16,14 @@
     pass
 
 
+with mock.patch.object(
+    self.my_runner, "first_method", autospec=True
+) as mock_run_adb, mock.patch.object(
+    self.my_runner, "second_method", autospec=True, return_value="foo"
+):
+    pass
+
+
 # output
 # This file doesn't use any Python 3.9+ only grammars.
 
@@ -28,3 +36,11 @@
 
 with make_context_manager1() as cm1, make_context_manager2() as cm2, make_context_manager3() as cm3, make_context_manager4() as cm4:
     pass
+
+
+with mock.patch.object(
+    self.my_runner, "first_method", autospec=True
+) as mock_run_adb, mock.patch.object(
+    self.my_runner, "second_method", autospec=True, return_value="foo"
+):
+    pass

tests/data/preview_context_managers/targeting_py38.py

@@ -23,6 +23,14 @@
     pass
 
 
+with mock.patch.object(
+    self.my_runner, "first_method", autospec=True
+) as mock_run_adb, mock.patch.object(
+    self.my_runner, "second_method", autospec=True, return_value="foo"
+):
+    pass
+
+
 # output
 
 
@@ -36,3 +44,11 @@
 
 with new_new_new1() as cm1, new_new_new2():
     pass
+
+
+with mock.patch.object(
+    self.my_runner, "first_method", autospec=True
+) as mock_run_adb, mock.patch.object(
+    self.my_runner, "second_method", autospec=True, return_value="foo"
+):
+    pass

tests/data/preview_context_managers/targeting_py39.py

@@ -49,6 +49,24 @@
     pass
 
 
+with mock.patch.object(
+    self.my_runner, "first_method", autospec=True
+) as mock_run_adb, mock.patch.object(
+    self.my_runner, "second_method", autospec=True, return_value="foo"
+):
+    pass
+
+
+with xxxxxxxx.some_kind_of_method(
+    some_argument=[
+        "first",
+        "second",
+        "third",
+    ]
+).another_method() as cmd:
+    pass
+
+
 # output
 
 
@@ -102,3 +120,22 @@
     ) as cm2,
 ):
     pass
+
+
+with (
+    mock.patch.object(self.my_runner, "first_method", autospec=True) as mock_run_adb,
+    mock.patch.object(
+        self.my_runner, "second_method", autospec=True, return_value="foo"
+    ),
+):
+    pass
+
+
+with xxxxxxxx.some_kind_of_method(
+    some_argument=[
+        "first",
+        "second",
+        "third",
+    ]
+).another_method() as cmd:
+    pass