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

Fixes #3472

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
 

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
 

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,

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

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