From db30456916ca733d011a2774e66beb3dd7064e14 Mon Sep 17 00:00:00 2001
From: Richard Si <63936253+ichard26@users.noreply.github.com>
Date: Mon, 26 Apr 2021 11:28:42 -0400
Subject: [PATCH] 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                          |  3 ++
 src/black/__init__.py               |  7 +++-
 tests/data/pep_572_remove_parens.py | 52 +++++++++++++++++++++++++++++
 3 files changed, 61 insertions(+), 1 deletion(-)

diff --git a/CHANGES.md b/CHANGES.md
index aa08396..9fe6b7e 100644
--- a/CHANGES.md
+++ b/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_
diff --git a/src/black/__init__.py b/src/black/__init__.py
index 00d3729..99afc7d 100644
--- a/src/black/__init__.py
+++ b/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]
diff --git a/tests/data/pep_572_remove_parens.py b/tests/data/pep_572_remove_parens.py
index 04cc75b..9a80485 100644
--- a/tests/data/pep_572_remove_parens.py
+++ b/tests/data/pep_572_remove_parens.py
@@ -17,6 +17,32 @@ 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)
+
 # output
 if foo := 0:
     pass
@@ -36,3 +62,29 @@ y += (x := 0)
 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)
-- 
2.39.5