Fix crash with walrus + await + with

Fixes psf#3472
JelleZijlstra · Dec 25, 2022 · c31425e · c31425e
1 parent 3feff21
commit c31425e
Show file tree

Hide file tree

Showing 3 changed files with 9 additions and 0 deletions.
diff --git a/CHANGES.md b/CHANGES.md
@@ -33,6 +33,7 @@
   string lambda values are now wrapped in parentheses (#3440)
 - Exclude string type annotations from improved string processing; fix crash when the
   return type annotation is stringified and spans across multiple lines (#3462)
+- Fix a crash in preview style with walrus operators used in `await` statements (#3473)
 
 ### Configuration
 

diff --git a/src/black/linegen.py b/src/black/linegen.py
@@ -1270,6 +1270,7 @@ 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,

diff --git a/tests/data/py_38/pep_572_remove_parens.py b/tests/data/py_38/pep_572_remove_parens.py
@@ -49,6 +49,9 @@ def a():
 def this_is_so_dumb() -> (please := no):
     pass
 
+async def await_the_walrus():
+    with (x := await y):
+        pass
 
 # output
 if foo := 0:
@@ -103,3 +106,7 @@ def a():
 def this_is_so_dumb() -> (please := no):
     pass
 
+
+async def await_the_walrus():
+    with (x := await y):
+        pass