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

Author: yilei

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)
 

src/black/linegen.py

@@ -36,6 +36,7 @@
     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 @@ def line(self, indent: int = 0) -> Iterator[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 @@ def visit_async_stmt(self, node: Node) -> Iterator[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."""

src/black/lines.py

@@ -28,7 +28,7 @@
     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 @@ def is_import(self) -> bool:
         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

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

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}
     )
 
 

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

tests/data/preview_context_managers/targeting_py39.py

@@ -67,6 +67,23 @@
     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 @@
     ]
 ).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