Don't strip parens in assert / return with assign expr

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)