From: Jelle Zijlstra <jelle.zijlstra@gmail.com>
Date: Wed, 18 Jan 2023 06:25:05 +0000 (-0800)
Subject: Fix crash with walrus + await + with (#3473)
X-Git-Url: https://git.madduck.net/etc/vim.git/commitdiff_plain/7e6d3fac197395b0a2b380cc60811536fe23626b?ds=sidebyside

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

Fixes #3472
---

diff --git a/CHANGES.md b/CHANGES.md
index 17dc0d6..97b68b9 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 da41886..14f8511 100644
--- 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_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 a11fb7c..a588077 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 20e80a6..05619dd 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 @@ with (y := [3, 2, 1]) as (funfunfun := indeed):
 @(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 9718d95..4e95fb0 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