From 7e6d3fac197395b0a2b380cc60811536fe23626b Mon Sep 17 00:00:00 2001 From: Jelle Zijlstra Date: Tue, 17 Jan 2023 22:25:05 -0800 Subject: [PATCH] Fix crash with walrus + await + with (#3473) Fixes #3472 --- CHANGES.md | 2 + src/black/linegen.py | 4 ++ src/black/nodes.py | 11 +++++ .../data/fast/pep_572_do_not_remove_parens.py | 4 ++ tests/data/py_38/pep_572_remove_parens.py | 40 +++++++++++++++++++ 5 files changed, 61 insertions(+) diff --git a/CHANGES.md b/CHANGES.md index 17dc0d686df..97b68b90fb1 100644 --- a/CHANGES.md +++ b/CHANGES.md @@ -35,6 +35,8 @@ - Fix two crashes in preview style involving edge cases with docstrings (#3451) - Exclude string type annotations from improved string processing; fix crash when the return type annotation is stringified and spans across multiple lines (#3462) +- Fix several crashes in preview style with walrus operators used in `with` statements + or tuples (#3473) ### Configuration diff --git a/src/black/linegen.py b/src/black/linegen.py index da41886f80d..14f851161fd 100644 --- a/src/black/linegen.py +++ b/src/black/linegen.py @@ -46,6 +46,7 @@ is_rpar_token, is_stub_body, is_stub_suite, + is_tuple_containing_walrus, is_vararg, is_walrus_assignment, is_yield, @@ -1279,6 +1280,7 @@ def maybe_make_parens_invisible_in_atom( not remove_brackets_around_comma and max_delimiter_priority_in_atom(node) >= COMMA_PRIORITY ) + or is_tuple_containing_walrus(node) ): return False @@ -1290,9 +1292,11 @@ def maybe_make_parens_invisible_in_atom( syms.return_stmt, syms.except_clause, syms.funcdef, + syms.with_stmt, # these ones aren't useful to end users, but they do please fuzzers syms.for_stmt, syms.del_stmt, + syms.for_stmt, ]: return False diff --git a/src/black/nodes.py b/src/black/nodes.py index a11fb7cc071..a588077f4de 100644 --- a/src/black/nodes.py +++ b/src/black/nodes.py @@ -563,6 +563,17 @@ def is_one_tuple(node: LN) -> bool: ) +def is_tuple_containing_walrus(node: LN) -> bool: + """Return True if `node` holds a tuple that contains a walrus operator.""" + if node.type != syms.atom: + return False + gexp = unwrap_singleton_parenthesis(node) + if gexp is None or gexp.type != syms.testlist_gexp: + return False + + return any(child.type == syms.namedexpr_test for child in gexp.children) + + def is_one_sequence_between( opening: Leaf, closing: Leaf, diff --git a/tests/data/fast/pep_572_do_not_remove_parens.py b/tests/data/fast/pep_572_do_not_remove_parens.py index 20e80a69377..05619ddcc2b 100644 --- a/tests/data/fast/pep_572_do_not_remove_parens.py +++ b/tests/data/fast/pep_572_do_not_remove_parens.py @@ -19,3 +19,7 @@ @(please := stop) def sigh(): pass + + +for (x := 3, y := 4) in y: + pass diff --git a/tests/data/py_38/pep_572_remove_parens.py b/tests/data/py_38/pep_572_remove_parens.py index 9718d95b499..4e95fb07f3a 100644 --- a/tests/data/py_38/pep_572_remove_parens.py +++ b/tests/data/py_38/pep_572_remove_parens.py @@ -49,6 +49,26 @@ def a(): def this_is_so_dumb() -> (please := no): pass +async def await_the_walrus(): + with (x := y): + pass + + with (x := y) as z, (a := b) as c: + pass + + with (x := await y): + pass + + with (x := await a, y := await b): + pass + + # Ideally we should remove one set of parentheses + with ((x := await a, y := await b)): + pass + + with (x := await a), (y := await b): + pass + # output if foo := 0: @@ -103,3 +123,23 @@ def a(): def this_is_so_dumb() -> (please := no): pass + +async def await_the_walrus(): + with (x := y): + pass + + with (x := y) as z, (a := b) as c: + pass + + with (x := await y): + pass + + with (x := await a, y := await b): + pass + + # Ideally we should remove one set of parentheses + with ((x := await a, y := await b)): + pass + + with (x := await a), (y := await b): + pass