Parenthesize with statements (#5758)

<!--
Thank you for contributing to Ruff! To help us out with reviewing,
please consider the following:

- Does this pull request include a summary of the change? (See below.)
- Does this pull request include a descriptive title?
- Does this pull request include references to any relevant issues?
-->

## Summary

This PR improves the parentheses handling for with items to get closer
to black's formatting.

### Case 1:

```python
# Black / Input
with (
    [
        "aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa",
        "bbbbbbbbbb",
        "cccccccccccccccccccccccccccccccccccccccccc",
        dddddddddddddddddddddddddddddddd,
    ] as example1,
    aaaaaaaaaaaaaaaaaaaaaaaaaa
    + bbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbb
    + cccccccccccccccccccccccccccc
    + ddddddddddddddddd as example2,
    CtxManager2() as example2,
    CtxManager2() as example2,
    CtxManager2() as example2,
):
    ...

# Before
with (
    [
        "aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa",
        "bbbbbbbbbb",
        "cccccccccccccccccccccccccccccccccccccccccc",
        dddddddddddddddddddddddddddddddd,
    ] as example1,
    (
        aaaaaaaaaaaaaaaaaaaaaaaaaa
        + bbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbb
        + cccccccccccccccccccccccccccc
        + ddddddddddddddddd
    ) as example2,
    CtxManager2() as example2,
    CtxManager2() as example2,
    CtxManager2() as example2,
):
    ...
```

Notice how Ruff wraps the binary expression in an extra set of
parentheses


### Case 2:
Black does not expand the with-items if the with has no parentheses:

```python
# Black / Input
with aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa + bbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbb as c:
    ...

# Before
with (
    aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa + bbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbb as c
):
    ...
```

Or 

```python
# Black / Input
with [
    "aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa",
    "bbbbbbbbbb",
    "cccccccccccccccccccccccccccccccccccccccccc",
    dddddddddddddddddddddddddddddddd,
] as example1, aaaaaaaaaaaaaaaaaaaaaaaaaa * bbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbb * cccccccccccccccccccccccccccc + ddddddddddddddddd as example2, CtxManager222222222222222() as example2:
    ...

# Before (Same as Case 1)
with (
    [
        "aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa",
        "bbbbbbbbbb",
        "cccccccccccccccccccccccccccccccccccccccccc",
        dddddddddddddddddddddddddddddddd,
    ] as example1,
    (
        aaaaaaaaaaaaaaaaaaaaaaaaaa
        * bbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbb
        * cccccccccccccccccccccccccccc
        + ddddddddddddddddd
    ) as example2,
    CtxManager222222222222222() as example2,
):
    ...

```
## Test Plan

I added new snapshot tests

Improves the django similarity index from 0.973 to 0.977

crates/ruff_python_formatter/src/expression/expr_await.rs

@@ -15,24 +15,27 @@ impl FormatNodeRule<ExprAwait> for FormatExprAwait {
     fn fmt_fields(&self, item: &ExprAwait, f: &mut PyFormatter) -> FormatResult<()> {
         let ExprAwait { range: _, value } = item;
 
-        let format_value = format_with(|f: &mut PyFormatter| {
-            if f.context().node_level().is_parenthesized() {
-                value.format().fmt(f)
-            } else {
-                maybe_parenthesize_expression(value, item, Parenthesize::Optional).fmt(f)
-            }
-        });
-
-        write!(f, [text("await"), space(), format_value])
+        write!(
+            f,
+            [
+                text("await"),
+                space(),
+                maybe_parenthesize_expression(value, item, Parenthesize::IfRequired)
+            ]
+        )
     }
 }
 
 impl NeedsParentheses for ExprAwait {
     fn needs_parentheses(
         &self,
-        _parent: AnyNodeRef,
+        parent: AnyNodeRef,
         _context: &PyFormatContext,
     ) -> OptionalParentheses {
-        OptionalParentheses::Multiline
+        if parent.is_expr_await() {
+            OptionalParentheses::Always
+        } else {
+            OptionalParentheses::Multiline
+        }
     }
 }

crates/ruff_python_formatter/src/expression/expr_bin_op.rs

@@ -176,9 +176,13 @@ impl FormatRule<Operator, PyFormatContext<'_>> for FormatOperator {
 impl NeedsParentheses for ExprBinOp {
     fn needs_parentheses(
         &self,
-        _parent: AnyNodeRef,
+        parent: AnyNodeRef,
         _context: &PyFormatContext,
     ) -> OptionalParentheses {
-        OptionalParentheses::Multiline
+        if parent.is_expr_await() && !self.op.is_pow() {
+            OptionalParentheses::Always
+        } else {
+            OptionalParentheses::Multiline
+        }
     }
 }

crates/ruff_python_formatter/src/expression/mod.rs

@@ -159,22 +159,31 @@ impl Format<PyFormatContext<'_>> for MaybeParenthesizeExpression<'_> {
             parenthesize,
         } = self;
 
-        let parenthesize = match parenthesize {
-            Parenthesize::Optional => {
-                is_expression_parenthesized(AnyNodeRef::from(*expression), f.context().source())
+        let comments = f.context().comments();
+        let preserve_parentheses = parenthesize.is_optional()
+            && is_expression_parenthesized(AnyNodeRef::from(*expression), f.context().source());
+
+        let has_comments = comments.has_leading_comments(*expression)
+            || comments.has_trailing_own_line_comments(*expression);
+
+        if preserve_parentheses || has_comments {
+            return expression.format().with_options(Parentheses::Always).fmt(f);
+        }
+
+        let needs_parentheses = expression.needs_parentheses(*parent, f.context());
+        let needs_parentheses = match parenthesize {
+            Parenthesize::IfRequired => {
+                if !needs_parentheses.is_always() && f.context().node_level().is_parenthesized() {
+                    OptionalParentheses::Never
+                } else {
+                    needs_parentheses
+                }
             }
-            Parenthesize::IfBreaks => false,
+            Parenthesize::Optional | Parenthesize::IfBreaks => needs_parentheses,
         };
 
-        let parentheses =
-            if parenthesize || f.context().comments().has_leading_comments(*expression) {
-                OptionalParentheses::Always
-            } else {
-                expression.needs_parentheses(*parent, f.context())
-            };
-
-        match parentheses {
-            OptionalParentheses::Multiline => {
+        match needs_parentheses {
+            OptionalParentheses::Multiline if *parenthesize != Parenthesize::IfRequired => {
                 if can_omit_optional_parentheses(expression, f.context()) {
                     optional_parentheses(&expression.format().with_options(Parentheses::Never))
                         .fmt(f)
@@ -186,7 +195,7 @@ impl Format<PyFormatContext<'_>> for MaybeParenthesizeExpression<'_> {
             OptionalParentheses::Always => {
                 expression.format().with_options(Parentheses::Always).fmt(f)
             }
-            OptionalParentheses::Never => {
+            OptionalParentheses::Never | OptionalParentheses::Multiline => {
                 expression.format().with_options(Parentheses::Never).fmt(f)
             }
         }

crates/ruff_python_formatter/src/expression/parentheses.rs

@@ -19,6 +19,12 @@ pub(crate) enum OptionalParentheses {
     Never,
 }
 
+impl OptionalParentheses {
+    pub(crate) const fn is_always(self) -> bool {
+        matches!(self, OptionalParentheses::Always)
+    }
+}
+
 pub(crate) trait NeedsParentheses {
     /// Determines if this object needs optional parentheses or if it is safe to omit the parentheses.
     fn needs_parentheses(
@@ -36,6 +42,17 @@ pub(crate) enum Parenthesize {
 
     /// Parenthesizes the expression only if it doesn't fit on a line.
     IfBreaks,
+
+    /// Only adds parentheses if absolutely necessary:
+    /// * The expression is not enclosed by another parenthesized expression and it expands over multiple lines
+    /// * The expression has leading or trailing comments. Adding parentheses is desired to prevent the comments from wandering.
+    IfRequired,
+}
+
+impl Parenthesize {
+    pub(crate) const fn is_optional(self) -> bool {
+        matches!(self, Parenthesize::Optional)
+    }
 }
 
 /// Whether it is necessary to add parentheses around an expression.