Don't strip parens in assert / return with assign expr (#2143)

Black would previously strip the parenthesis away from statements like this these ones:

    assert (spam := 12 + 1)
    return (cheese := 1 - 12)

Which happens to be invalid code. Now before making the parenthesis invisible, Black
checks if the assignment expression's parent is an assert stamtment, aborting if True.

Raise, yield, and await are already handled fine.

I added a bunch of test cases from the PEP defining asssignment expressions (PEP 572).

CHANGES.md

@@ -8,6 +8,9 @@
   in the name of the cache file. Without this fix, changes in these flags would not take
   effect if the cache had already been populated. (#2131)
 
+- Don't remove necessary parentheses from assignment expression containing assert /
+  return statements. (#2143)
+
 ### 21.4b0
 
 #### _Black_

src/black/__init__.py

@@ -5613,7 +5613,12 @@ def maybe_make_parens_invisible_in_atom(node: LN, parent: LN) -> bool:
         return False
 
     if is_walrus_assignment(node):
-        if parent.type in [syms.annassign, syms.expr_stmt]:
+        if parent.type in [
+            syms.annassign,
+            syms.expr_stmt,
+            syms.assert_stmt,
+            syms.return_stmt,
+        ]:
             return False
 
     first = node.children[0]

tests/data/pep_572_remove_parens.py

@@ -17,6 +17,32 @@
 
 a, b = (test := (1, 2))
 
+# see also https://github.com/psf/black/issues/2139
+assert (foo := 42 - 12)
+
+foo(x=(y := f(x)))
+
+
+def foo(answer=(p := 42)):
+    ...
+
+
+def foo2(answer: (p := 42) = 5):
+    ...
+
+
+lambda: (x := 1)
+
+# we don't touch expressions in f-strings but if we do one day, don't break 'em
+f'{(x:=10)}'
+
+
+def a():
+    return (x := 3)
+    await (b := 1)
+    yield (a := 2)
+    raise (c := 3)
+
 # output
 if foo := 0:
     pass
@@ -36,3 +62,29 @@
 test: int = (test2 := 2)
 
 a, b = (test := (1, 2))
+
+# see also https://github.com/psf/black/issues/2139
+assert (foo := 42 - 12)
+
+foo(x=(y := f(x)))
+
+
+def foo(answer=(p := 42)):
+    ...
+
+
+def foo2(answer: (p := 42) = 5):
+    ...
+
+
+lambda: (x := 1)
+
+# we don't touch expressions in f-strings but if we do one day, don't break 'em
+f"{(x:=10)}"
+
+
+def a():
+    return (x := 3)
+    await (b := 1)
+    yield (a := 2)
+    raise (c := 3)