From 7e6d3fac197395b0a2b380cc60811536fe23626b Mon Sep 17 00:00:00 2001
From: Jelle Zijlstra <jelle.zijlstra@gmail.com>
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