Improve single with item pre 3.9 formatting

astral-sh · Mar 8, 2024 · f8937a3 · f8937a3
1 parent a42c365
commit f8937a3
Show file tree

Hide file tree

Showing 5 changed files with 167 additions and 4 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
@@ -318,3 +318,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
@@ -104,7 +104,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(),
 

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
@@ -324,6 +324,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
@@ -688,6 +698,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
+@@ -49,7 +49,9 @@
+ with a:  # should remove brackets
+     pass
+ 
+-with aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa + bbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbb as c:
++with (
++    aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa + bbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbb
++) as c:
+     pass
+ 
+ 
+@@ -214,7 +216,9 @@
+     pass
+ 
+ # Breaking of with items.
+-with test as (  # bar  # foo
++with (
++    test  # bar
++) as (  # foo
+     # test
+     foo
+ ):
+@@ -226,7 +230,9 @@
+ ):
+     pass
+ 
+-with test as (  # bar  # foo  # baz
++with (
++    test  # bar
++) as (  # foo  # baz
+     # test
+     foo
+ ):
+@@ -279,7 +285,9 @@
+ ) as b:
+     pass
+ 
+-with aaaaaaaaaaaaaaaaaaaaaaaaaaaaaa + bbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbb as b:
++with (
++    aaaaaaaaaaaaaaaaaaaaaaaaaaaaaa + bbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbb
++) as b:
+     pass
+ 
+ with (
+@@ -322,15 +330,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(
+@@ -344,14 +356,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
 ```
 
 
@@ -1093,4 +1218,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
 ```