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).

diff --git a/CHANGES.md b/CHANGES.md
index aa08396e49196111a5876e926c4a96ae45a7b51a..9fe6b7e3fa4afd403ef25879267b7b745267be27 100644 (file)
--- 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 00d3729c8354496812f9250d1f6f6f5530572cbe..99afc7d3c507fe06d0f21a1f1f50ea4a4b0577ad 100644 (file)
--- 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 04cc75bc36c449a6d047832f3265fa70e2a3f6d6..9a804859a09b92f321603004d79634044afb859b 100644 (file)
--- 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)