From fc6cea0f0e0f6da99fffbedbfcf3d50cc0e641a4 Mon Sep 17 00:00:00 2001 From: "Yilei \"Dolee\" Yang" Date: Thu, 16 Mar 2023 13:31:27 -0700 Subject: [PATCH] Consistently format async statements similar to their non-async version. (#3609) --- CHANGES.md | 2 ++ src/black/linegen.py | 19 +++++++++-- src/black/lines.py | 8 ++--- src/black/mode.py | 1 + src/black/nodes.py | 21 ++++++++++-- tests/data/preview/async_stmts.py | 27 +++++++++++++++ .../targeting_py39.py | 33 +++++++++++++++++++ 7 files changed, 103 insertions(+), 8 deletions(-) create mode 100644 tests/data/preview/async_stmts.py diff --git a/CHANGES.md b/CHANGES.md index 2fa0cb4..eff2640 100644 --- 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 6f67799..b6b83da 100644 --- 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 4b57d1f..329dfc4 100644 --- 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 d703889..6af0417 100644 --- 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 90728f3..4e9411b 100644 --- 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 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 643c6fd..c9fcf9c 100644 --- 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 -- 2.39.5