Consistently format async statements similar to their non-async version. (#3609)

diff --git a/CHANGES.md b/CHANGES.md
index 2fa0cb41b38388af9be2dba8aab43190666c9f27..eff2640a01e9f3191389b3f7d3ae5a99e7571451 100644 (file)
--- a/CHANGES.md
+++ b/CHANGES.md
@@ -16,6 +16,8 @@
 
 - Add trailing commas to collection literals even if there's a comment after the last
   entry (#3393)
+- `async def`, `async for`, and `async with` statements are now formatted consistently
+  compared to their non-async version. (#3609)
 - `with` statements that contain two context managers will be consistently wrapped in
   parentheses (#3589)
 

diff --git a/src/black/linegen.py b/src/black/linegen.py
index 6f67799e7172e033fdf5ffd8e4eb3473b64be943..b6b83da26f776fa4810c18f42a1c911bc2eea67e 100644 (file)
--- a/src/black/linegen.py
+++ b/src/black/linegen.py
@@ -36,6 +36,7 @@ from black.nodes import (
     Visitor,
     ensure_visible,
     is_arith_like,
+    is_async_stmt_or_funcdef,
     is_atom_with_invisible_parens,
     is_docstring,
     is_empty_tuple,
@@ -110,6 +111,17 @@ class LineGenerator(Visitor[Line]):
             self.current_line.depth += indent
             return  # Line is empty, don't emit. Creating a new one unnecessary.
 
+        if (
+            Preview.improved_async_statements_handling in self.mode
+            and len(self.current_line.leaves) == 1
+            and is_async_stmt_or_funcdef(self.current_line.leaves[0])
+        ):
+            # Special case for async def/for/with statements. `visit_async_stmt`
+            # adds an `ASYNC` leaf then visits the child def/for/with statement
+            # nodes. Line yields from those nodes shouldn't treat the former
+            # `ASYNC` leaf as a complete line.
+            return
+
         complete_line = self.current_line
         self.current_line = Line(mode=self.mode, depth=complete_line.depth + indent)
         yield complete_line
@@ -301,8 +313,11 @@ class LineGenerator(Visitor[Line]):
                 break
 
         internal_stmt = next(children)
-        for child in internal_stmt.children:
-            yield from self.visit(child)
+        if Preview.improved_async_statements_handling in self.mode:
+            yield from self.visit(internal_stmt)
+        else:
+            for child in internal_stmt.children:
+                yield from self.visit(child)
 
     def visit_decorators(self, node: Node) -> Iterator[Line]:
         """Visit decorators."""

diff --git a/src/black/lines.py b/src/black/lines.py
index 4b57d1f0ea8c3db10328e52ae94f6c88c9df3ff0..329dfc4f0d3d7de0dbe545656c68d020628a9a5a 100644 (file)
--- a/src/black/lines.py
+++ b/src/black/lines.py
@@ -28,7 +28,7 @@ from black.nodes import (
     is_multiline_string,
     is_one_sequence_between,
     is_type_comment,
-    is_with_stmt,
+    is_with_or_async_with_stmt,
     replace_child,
     syms,
     whitespace,
@@ -124,9 +124,9 @@ class Line:
         return bool(self) and is_import(self.leaves[0])
 
     @property
-    def is_with_stmt(self) -> bool:
+    def is_with_or_async_with_stmt(self) -> bool:
         """Is this a with_stmt line?"""
-        return bool(self) and is_with_stmt(self.leaves[0])
+        return bool(self) and is_with_or_async_with_stmt(self.leaves[0])
 
     @property
     def is_class(self) -> bool:
@@ -872,7 +872,7 @@ def can_omit_invisible_parens(
         if (
             Preview.wrap_multiple_context_managers_in_parens in line.mode
             and max_priority == COMMA_PRIORITY
-            and rhs.head.is_with_stmt
+            and rhs.head.is_with_or_async_with_stmt
         ):
             # For two context manager with statements, the optional parentheses read
             # better. In this case, `rhs.body` is the context managers part of

diff --git a/src/black/mode.py b/src/black/mode.py
index d70388916dabf1053d62f20921cb83f64abd9922..6af041791937a0f6e29ff77f82bded006f5bbda1 100644 (file)
--- a/src/black/mode.py
+++ b/src/black/mode.py
@@ -155,6 +155,7 @@ class Preview(Enum):
 
     add_trailing_comma_consistently = auto()
     hex_codes_in_unicode_sequences = auto()
+    improved_async_statements_handling = auto()
     multiline_string_handling = auto()
     prefer_splitting_right_hand_side_of_assignments = auto()
     # NOTE: string_processing requires wrap_long_dict_values_in_parens

diff --git a/src/black/nodes.py b/src/black/nodes.py
index 90728f3c87c2e60e098f65ba905e4528d39052e0..4e9411b1b7988066f3e605504b9b70a6a10d8d3b 100644 (file)
--- a/src/black/nodes.py
+++ b/src/black/nodes.py
@@ -789,13 +789,30 @@ def is_import(leaf: Leaf) -> bool:
     )
 
 
-def is_with_stmt(leaf: Leaf) -> bool:
-    """Return True if the given leaf starts a with statement."""
+def is_with_or_async_with_stmt(leaf: Leaf) -> bool:
+    """Return True if the given leaf starts a with or async with statement."""
     return bool(
         leaf.type == token.NAME
         and leaf.value == "with"
         and leaf.parent
         and leaf.parent.type == syms.with_stmt
+    ) or bool(
+        leaf.type == token.ASYNC
+        and leaf.next_sibling
+        and leaf.next_sibling.type == syms.with_stmt
+    )
+
+
+def is_async_stmt_or_funcdef(leaf: Leaf) -> bool:
+    """Return True if the given leaf starts an async def/for/with statement.
+
+    Note that `async def` can be either an `async_stmt` or `async_funcdef`,
+    the latter is used when it has decorators.
+    """
+    return bool(
+        leaf.type == token.ASYNC
+        and leaf.parent
+        and leaf.parent.type in {syms.async_stmt, syms.async_funcdef}
     )
 
 

diff --git a/tests/data/preview/async_stmts.py b/tests/data/preview/async_stmts.py
new file mode 100644 (file)
index 0000000..fe9594b
--- /dev/null
+++ b/tests/data/preview/async_stmts.py
@@ -0,0 +1,27 @@
+async def func() -> (int):
+    return 0
+
+
+@decorated
+async def func() -> (int):
+    return 0
+
+
+async for (item) in async_iter:
+    pass
+
+
+# output
+
+
+async def func() -> int:
+    return 0
+
+
+@decorated
+async def func() -> int:
+    return 0
+
+
+async for item in async_iter:
+    pass

diff --git a/tests/data/preview_context_managers/targeting_py39.py b/tests/data/preview_context_managers/targeting_py39.py
index 643c6fd958b3e9cde27d04533131144ba8516729..c9fcf9c8ba2c9699382bfb45f57e75c0b344d98a 100644 (file)
--- a/tests/data/preview_context_managers/targeting_py39.py
+++ b/tests/data/preview_context_managers/targeting_py39.py
@@ -67,6 +67,23 @@ with xxxxxxxx.some_kind_of_method(
     pass
 
 
+async def func():
+    async with \
+        make_context_manager1() as cm1, \
+        make_context_manager2() as cm2, \
+        make_context_manager3() as cm3, \
+        make_context_manager4() as cm4 \
+    :
+        pass
+
+    async with some_function(
+        argument1, argument2, argument3="some_value"
+    ) as some_cm, some_other_function(
+        argument1, argument2, argument3="some_value"
+    ):
+        pass
+
+
 # output
 
 
@@ -139,3 +156,19 @@ with xxxxxxxx.some_kind_of_method(
     ]
 ).another_method() as cmd:
     pass
+
+
+async def func():
+    async with (
+        make_context_manager1() as cm1,
+        make_context_manager2() as cm2,
+        make_context_manager3() as cm3,
+        make_context_manager4() as cm4,
+    ):
+        pass
+
+    async with (
+        some_function(argument1, argument2, argument3="some_value") as some_cm,
+        some_other_function(argument1, argument2, argument3="some_value"),
+    ):
+        pass