Fix crash with walrus + await + with (#3473)

Fixes #3472

diff --git a/CHANGES.md b/CHANGES.md
index 17dc0d686dfa19d93346e4d6a5b5c2d8360432a6..97b68b90fb18d2ac2e6023f8cbf72084faf83667 100644 (file)
--- 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 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
 
 
 ### Configuration
 

diff --git a/src/black/linegen.py b/src/black/linegen.py
index da41886f80d616845db26eadc5eac6a481cce7e1..14f851161fd60cbee9d7cf0cf058d03f469a5461 100644 (file)
--- a/src/black/linegen.py
+++ b/src/black/linegen.py
@@ -46,6 +46,7 @@ from black.nodes import (
     is_rpar_token,
     is_stub_body,
     is_stub_suite,
     is_rpar_token,
     is_stub_body,
     is_stub_suite,

+    is_tuple_containing_walrus,

     is_vararg,
     is_walrus_assignment,
     is_yield,
     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
         )
             not remove_brackets_around_comma
             and max_delimiter_priority_in_atom(node) >= COMMA_PRIORITY
         )

+        or is_tuple_containing_walrus(node)

     ):
         return False
 
     ):
         return False
 


@@ -1290,9 +1292,11 @@ def maybe_make_parens_invisible_in_atom(
             syms.return_stmt,
             syms.except_clause,
             syms.funcdef,
             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,
             # these ones aren't useful to end users, but they do please fuzzers
             syms.for_stmt,
             syms.del_stmt,

+            syms.for_stmt,

         ]:
             return False
 
         ]:
             return False
 

diff --git a/src/black/nodes.py b/src/black/nodes.py
index a11fb7cc071647fe58cdd7eaac516afa47088e57..a588077f4def9eb028da74463119f86f57d125f9 100644 (file)
--- 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,
 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 20e80a693774a50efb6f969fe6de3f3eeda6370f..05619ddcc2b0336560560d5adaac81d100a8d82a 100644 (file)
--- 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 @@ with (y := [3, 2, 1]) as (funfunfun := indeed):
 @(please := stop)
 def sigh():
     pass
 @(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 9718d95b499de4e2c966de31d10433deb471e681..4e95fb07f3af5bdd0016b98b8b81267bfada6310 100644 (file)
--- 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
 
 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:
 
 # output
 if foo := 0:


@@ -103,3 +123,23 @@ def a():
 def this_is_so_dumb() -> (please := no):
     pass
 
 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