From dca430f4d212efa44b3d76ec43736e7dc7a9500d Mon Sep 17 00:00:00 2001
From: konsti <konstin@mailbox.org>
Date: Fri, 17 Nov 2023 18:24:19 +0100
Subject: [PATCH] Fix instability with await fluent style (#8676)

Fix an instability where await was followed by a breaking fluent style
expression:

```python
test_data = await (
    Stream.from_async(async_data)
    .flat_map_async()
    .map()
    .filter_async(is_valid_data)
    .to_list()
)
```

Note that this technically a minor style change (see ecosystem check)
---
 .../test/fixtures/ruff/expression/await.py     |  9 +++++++++
 .../src/expression/expr_await.rs               |  2 +-
 .../src/expression/parentheses.rs              |  4 +++-
 ...y@simple_cases__remove_await_parens.py.snap | 13 +++++++++++--
 .../snapshots/format@expression__await.py.snap | 18 ++++++++++++++++++
 5 files changed, 42 insertions(+), 4 deletions(-)

diff --git a/crates/ruff_python_formatter/resources/test/fixtures/ruff/expression/await.py b/crates/ruff_python_formatter/resources/test/fixtures/ruff/expression/await.py
index 03b7ce38de033..28b6572fc48f0 100644
--- a/crates/ruff_python_formatter/resources/test/fixtures/ruff/expression/await.py
+++ b/crates/ruff_python_formatter/resources/test/fixtures/ruff/expression/await.py
@@ -48,3 +48,12 @@
     # comment
     [foo]
 )
+
+# https://github.com/astral-sh/ruff/issues/8644
+test_data = await (
+    Stream.from_async(async_data)
+    .flat_map_async()
+    .map()
+    .filter_async(is_valid_data)
+    .to_list()
+)
diff --git a/crates/ruff_python_formatter/src/expression/expr_await.rs b/crates/ruff_python_formatter/src/expression/expr_await.rs
index 26521bb3336b2..a7073312988c7 100644
--- a/crates/ruff_python_formatter/src/expression/expr_await.rs
+++ b/crates/ruff_python_formatter/src/expression/expr_await.rs
@@ -20,7 +20,7 @@ impl FormatNodeRule<ExprAwait> for FormatExprAwait {
             [
                 token("await"),
                 space(),
-                maybe_parenthesize_expression(value, item, Parenthesize::IfRequired)
+                maybe_parenthesize_expression(value, item, Parenthesize::IfBreaks)
             ]
         )
     }
diff --git a/crates/ruff_python_formatter/src/expression/parentheses.rs b/crates/ruff_python_formatter/src/expression/parentheses.rs
index 75e0ee3fe0def..d05b9dcbd7d5f 100644
--- a/crates/ruff_python_formatter/src/expression/parentheses.rs
+++ b/crates/ruff_python_formatter/src/expression/parentheses.rs
@@ -14,6 +14,7 @@ use crate::comments::{
 use crate::context::{NodeLevel, WithNodeLevel};
 use crate::prelude::*;
 
+/// From the perspective of the expression, under which circumstances does it need parentheses
 #[derive(Copy, Clone, Debug, Eq, PartialEq)]
 pub(crate) enum OptionalParentheses {
     /// Add parentheses if the expression expands over multiple lines
@@ -41,7 +42,8 @@ pub(crate) trait NeedsParentheses {
     ) -> OptionalParentheses;
 }
 
-/// Configures if the expression should be parenthesized.
+/// From the perspective of the parent statement or expression, when should the child expression
+/// get parentheses?
 #[derive(Copy, Clone, Debug, PartialEq)]
 pub(crate) enum Parenthesize {
     /// Parenthesizes the expression if it doesn't fit on a line OR if the expression is parenthesized in the source code.
diff --git a/crates/ruff_python_formatter/tests/snapshots/black_compatibility@simple_cases__remove_await_parens.py.snap b/crates/ruff_python_formatter/tests/snapshots/black_compatibility@simple_cases__remove_await_parens.py.snap
index a72df1daba4cf..33eabc6fa9b93 100644
--- a/crates/ruff_python_formatter/tests/snapshots/black_compatibility@simple_cases__remove_await_parens.py.snap
+++ b/crates/ruff_python_formatter/tests/snapshots/black_compatibility@simple_cases__remove_await_parens.py.snap
@@ -93,7 +93,7 @@ async def main():
 ```diff
 --- Black
 +++ Ruff
-@@ -21,7 +21,9 @@
+@@ -21,11 +21,15 @@
  
  # Check comments
  async def main():
@@ -103,6 +103,13 @@ async def main():
 +    )
  
  
+ async def main():
+-    await asyncio.sleep(1)  # Hello
++    await (
++        asyncio.sleep(1)  # Hello
++    )
+ 
+ 
  async def main():
 ```
 
@@ -138,7 +145,9 @@ async def main():
 
 
 async def main():
-    await asyncio.sleep(1)  # Hello
+    await (
+        asyncio.sleep(1)  # Hello
+    )
 
 
 async def main():
diff --git a/crates/ruff_python_formatter/tests/snapshots/format@expression__await.py.snap b/crates/ruff_python_formatter/tests/snapshots/format@expression__await.py.snap
index fa8d86bf521f5..9a9a7da7cbeb8 100644
--- a/crates/ruff_python_formatter/tests/snapshots/format@expression__await.py.snap
+++ b/crates/ruff_python_formatter/tests/snapshots/format@expression__await.py.snap
@@ -54,6 +54,15 @@ await (
     # comment
     [foo]
 )
+
+# https://github.com/astral-sh/ruff/issues/8644
+test_data = await (
+    Stream.from_async(async_data)
+    .flat_map_async()
+    .map()
+    .filter_async(is_valid_data)
+    .to_list()
+)
 ```
 
 ## Output
@@ -122,6 +131,15 @@ await (
     # comment
     [foo]
 )
+
+# https://github.com/astral-sh/ruff/issues/8644
+test_data = await (
+    Stream.from_async(async_data)
+    .flat_map_async()
+    .map()
+    .filter_async(is_valid_data)
+    .to_list()
+)
 ```