Treat walruses like other binary operators in subscripts (psf#4109)

cobaltt7 · web-flow · commit c80685f36183 · 2023-12-27T22:24:25.000-08:00
Fixes psf#4078
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,15 @@ def is_complex_subscript(self, leaf: Leaf) -> bool:
 
             if subscript_start.type == syms.subscriptlist:
                 subscript_start = child_towards(subscript_start, leaf)
+
+        # When this is moved out of preview, add syms.namedexpr_test directly to
+        # TEST_DESCENDANTS in nodes.py
+        if Preview.walrus_subscript in self.mode:
+            test_decendants = TEST_DESCENDANTS | {syms.namedexpr_test}
+        else:
+            test_decendants = TEST_DESCENDANTS
         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)]
