fix: In subscripts, treat walruses just like other binary operators

Signed-off-by: RedGuy12 <61329810+RedGuy12@users.noreply.github.com>
cobaltt7 · Dec 13, 2023 · 578f33e · 578f33e
1 parent 35ce37d
commit 578f33e
Show file tree

Hide file tree

Showing 3 changed files with 10 additions and 3 deletions.
diff --git a/CHANGES.md b/CHANGES.md
@@ -14,6 +14,9 @@
 
 <!-- Changes that affect Black's preview style -->
 
+- Fix bug where spaces were not added around parenthesized walruses in subscripts,
+  unlike other binary operators (#4109)
+
 ### Configuration
 
 <!-- Changes to how Black can be configured -->

diff --git a/src/black/lines.py b/src/black/lines.py
@@ -446,8 +446,12 @@ def is_complex_subscript(self, leaf: Leaf) -> bool:
 
             if subscript_start.type == syms.subscriptlist:
                 subscript_start = child_towards(subscript_start, leaf)
+
+        test_decendants = TEST_DESCENDANTS
+        if Preview.walrus_subscript in self.mode:
+            test_decendants.add(syms.namedexpr_test)
         return subscript_start is not None and any(
-            n.type in TEST_DESCENDANTS for n in subscript_start.pre_order()
+            n.type in test_decendants for n in subscript_start.pre_order()
         )
 
     def enumerate_with_length(

diff --git a/tests/data/cases/preview_pep_572.py b/tests/data/cases/preview_pep_572.py
@@ -3,5 +3,5 @@
 x[:(a:=0)]
 
 # output
-x[(a := 0):]
-x[:(a := 0)]
+x[(a := 0) :]
+x[: (a := 0)]