Improve single with item pre 3.9 formatting

astral-sh · Mar 7, 2024 · 4932a2d · 4932a2d
1 parent 25ada0a
commit 4932a2d
Show file tree

Hide file tree

Showing 5 changed files with 170 additions and 7 deletions.
diff --git a/crates/ruff_python_formatter/resources/test/fixtures/ruff/statement/with.py b/crates/ruff_python_formatter/resources/test/fixtures/ruff/statement/with.py
@@ -313,3 +313,13 @@
     )
 ):
     pass
+
+with aaaaaaaaaaaaaaaaaaaaaaaaaaaaaa + bbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbb:
+    pass
+
+with aaaaaaaaaaaaaaaaaaaaaaaaaaaaaa + bbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbb as b:
+    pass
+
+if True:
+    with anyio.CancelScope(shield=True) if get_running_loop() else contextlib.nullcontext() as b:
+        pass
diff --git a/crates/ruff_python_formatter/src/other/with_item.rs b/crates/ruff_python_formatter/src/other/with_item.rs
@@ -7,6 +7,7 @@ use crate::expression::parentheses::{
     is_expression_parenthesized, parenthesized, Parentheses, Parenthesize,
 };
 use crate::prelude::*;
+use crate::preview::is_with_single_item_pre_39_enabled;
 
 #[derive(Copy, Clone, Debug, Eq, PartialEq, Default)]
 pub enum WithItemLayout {
@@ -49,7 +50,7 @@ pub enum WithItemLayout {
     /// with a, b:
     ///     ...
     /// ```
-    Python38OrOlder,
+    Python38OrOlder { single: bool },
 
     /// A with item where the `with` formatting adds parentheses around all context managers if necessary.
     ///
@@ -135,8 +136,10 @@ impl FormatNodeRule<WithItem> for FormatWithItem {
                 )?;
             }
 
-            WithItemLayout::Python38OrOlder => {
-                let parenthesize = if is_parenthesized {
+            WithItemLayout::Python38OrOlder { single } => {
+                let parenthesize = if (single && is_with_single_item_pre_39_enabled(f.context()))
+                    || is_parenthesized
+                {
                     Parenthesize::IfBreaks
                 } else {
                     Parenthesize::IfRequired

diff --git a/crates/ruff_python_formatter/src/preview.rs b/crates/ruff_python_formatter/src/preview.rs
@@ -18,3 +18,7 @@ pub(crate) const fn is_hug_parens_with_braces_and_square_brackets_enabled(
 pub(crate) fn is_f_string_formatting_enabled(context: &PyFormatContext) -> bool {
     context.is_preview()
 }
+
+pub(crate) fn is_with_single_item_pre_39_enabled(context: &PyFormatContext) -> bool {
+    context.is_preview()
+}
diff --git a/crates/ruff_python_formatter/src/statement/stmt_with.rs b/crates/ruff_python_formatter/src/statement/stmt_with.rs
@@ -100,7 +100,9 @@ impl FormatNodeRule<StmtWith> for FormatStmtWith {
                             WithItemsLayout::Python38OrOlder => f
                                 .join_with(format_args![token(","), space()])
                                 .entries(with_stmt.items.iter().map(|item| {
-                                    item.format().with_options(WithItemLayout::Python38OrOlder)
+                                    item.format().with_options(WithItemLayout::Python38OrOlder {
+                                        single: with_stmt.items.len() == 1,
+                                    })
                                 }))
                                 .finish(),
 
@@ -255,6 +257,9 @@ impl<'a> WithItemsLayout<'a> {
             return Ok(Self::ParenthesizeIfExpands);
         }
 
+        let can_parenthesize = context.options().target_version() >= PythonVersion::Py39
+            || are_with_items_parenthesized(with, context)?;
+
         if let [single] = with.items.as_slice() {
             // If the with item itself has comments (not the context expression), then keep the parentheses
             // ```python
@@ -278,9 +283,6 @@ impl<'a> WithItemsLayout<'a> {
             }
         }
 
-        let can_parenthesize = context.options().target_version() >= PythonVersion::Py39
-            || are_with_items_parenthesized(with, context)?;
-
         // If the target version doesn't support parenthesized context managers and they aren't
         // parenthesized by the user, bail out.
         if !can_parenthesize {

diff --git a/crates/ruff_python_formatter/tests/snapshots/format@statement__with.py.snap b/crates/ruff_python_formatter/tests/snapshots/format@statement__with.py.snap
@@ -319,6 +319,16 @@ with (
     )
 ):
     pass
+
+with aaaaaaaaaaaaaaaaaaaaaaaaaaaaaa + bbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbb:
+    pass
+
+with aaaaaaaaaaaaaaaaaaaaaaaaaaaaaa + bbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbb as b:
+    pass
+
+if True:
+    with anyio.CancelScope(shield=True) if get_running_loop() else contextlib.nullcontext() as b:
+        pass
 ```
 
 ## Outputs
@@ -678,6 +688,121 @@ with open(
     "/etc/hosts"  # This is an incredibly long comment that has been replaced for sanitization
 ):
     pass
+
+with aaaaaaaaaaaaaaaaaaaaaaaaaaaaaa + bbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbb:
+    pass
+
+with aaaaaaaaaaaaaaaaaaaaaaaaaaaaaa + bbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbb as b:
+    pass
+
+if True:
+    with anyio.CancelScope(
+        shield=True
+    ) if get_running_loop() else contextlib.nullcontext() as b:
+        pass
+```
+
+
+#### Preview changes
+```diff
+--- Stable
++++ Preview
+@@ -45,7 +45,9 @@
+ with a:  # should remove brackets
+     pass
+ 
+-with aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa + bbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbb as c:
++with (
++    aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa + bbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbb
++) as c:
+     pass
+ 
+ # currently unparsable by black: https://github.com/psf/black/issues/3678
+@@ -209,7 +211,9 @@
+     pass
+ 
+ # Breaking of with items.
+-with test as (  # bar  # foo
++with (
++    test  # bar
++) as (  # foo
+     # test
+     foo
+ ):
+@@ -221,7 +225,9 @@
+ ):
+     pass
+ 
+-with test as (  # bar  # foo  # baz
++with (
++    test  # bar
++) as (  # foo  # baz
+     # test
+     foo
+ ):
+@@ -274,7 +280,9 @@
+ ) as b:
+     pass
+ 
+-with aaaaaaaaaaaaaaaaaaaaaaaaaaaaaa + bbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbb as b:
++with (
++    aaaaaaaaaaaaaaaaaaaaaaaaaaaaaa + bbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbb
++) as b:
+     pass
+ 
+ with (
+@@ -317,15 +325,19 @@
+         pass
+ 
+ if True:
+-    with anyio.CancelScope(
+-        shield=True
+-    ) if get_running_loop() else contextlib.nullcontext():
++    with (
++        anyio.CancelScope(shield=True)
++        if get_running_loop()
++        else contextlib.nullcontext()
++    ):
+         pass
+ 
+ if True:
+-    with anyio.CancelScope(
+-        shield=True
+-    ) if get_running_loop() else contextlib.nullcontext() as c:
++    with (
++        anyio.CancelScope(shield=True)
++        if get_running_loop()
++        else contextlib.nullcontext()
++    ) as c:
+         pass
+ 
+ with Child(aaaaaaaaa, bbbbbbbbbbbbbbb, cccccc), Document(
+@@ -339,14 +351,20 @@
+ ):
+     pass
+ 
+-with aaaaaaaaaaaaaaaaaaaaaaaaaaaaaa + bbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbb:
++with (
++    aaaaaaaaaaaaaaaaaaaaaaaaaaaaaa + bbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbb
++):
+     pass
+ 
+-with aaaaaaaaaaaaaaaaaaaaaaaaaaaaaa + bbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbb as b:
++with (
++    aaaaaaaaaaaaaaaaaaaaaaaaaaaaaa + bbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbb
++) as b:
+     pass
+ 
+ if True:
+-    with anyio.CancelScope(
+-        shield=True
+-    ) if get_running_loop() else contextlib.nullcontext() as b:
++    with (
++        anyio.CancelScope(shield=True)
++        if get_running_loop()
++        else contextlib.nullcontext()
++    ) as b:
+         pass
 ```
 
 
@@ -1078,4 +1203,23 @@ with open(
     "/etc/hosts"  # This is an incredibly long comment that has been replaced for sanitization
 ):
     pass
+
+with (
+    aaaaaaaaaaaaaaaaaaaaaaaaaaaaaa + bbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbb
+):
+    pass
+
+with (
+    aaaaaaaaaaaaaaaaaaaaaaaaaaaaaa
+    + bbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbb as b
+):
+    pass
+
+if True:
+    with (
+        anyio.CancelScope(shield=True)
+        if get_running_loop()
+        else contextlib.nullcontext() as b
+    ):
+        pass
 ```