Fix crash in return annotation with walrus

Fixes psf#3419
JelleZijlstra · Dec 10, 2022 · bae256b · bae256b
1 parent 96e62c5
commit bae256b
Show file tree

Hide file tree

Showing 3 changed files with 11 additions and 0 deletions.
diff --git a/CHANGES.md b/CHANGES.md
@@ -15,6 +15,7 @@
 <!-- Changes that affect Black's preview style -->
 
 - Fix a crash in preview style with assert + parenthesized string (#3415)
+- Fix a crash in preview style with a walrus operator in a return annotation (#3422)
 
 ### Configuration
 

diff --git a/src/black/linegen.py b/src/black/linegen.py
@@ -1160,6 +1160,7 @@ def maybe_make_parens_invisible_in_atom(
             syms.expr_stmt,
             syms.assert_stmt,
             syms.return_stmt,
+            syms.funcdef,
             # these ones aren't useful to end users, but they do please fuzzers
             syms.for_stmt,
             syms.del_stmt,

diff --git a/tests/data/preview/return_annotation_brackets.py b/tests/data/preview/return_annotation_brackets.py
@@ -91,6 +91,10 @@ def foo() -> tuple[int, int, int,]:
 def frobnicate() -> "ThisIsTrulyUnreasonablyExtremelyLongClassName | list[ThisIsTrulyUnreasonablyExtremelyLongClassName]":
     pass
 
+# Should not remove parens around walrus operator
+def walrus() -> (x := 1):
+    pass
+
 # output
 # Control
 def double(a: int) -> int:
@@ -220,3 +224,8 @@ def frobnicate() -> (
     " list[ThisIsTrulyUnreasonablyExtremelyLongClassName]"
 ):
     pass
+
+
+# Should not remove parens around walrus operator
+def walrus() -> (x := 1):
+    pass