From c80685f36183f146f831a5737510cf105f947745 Mon Sep 17 00:00:00 2001
From: cobalt <61329810+RedGuy12@users.noreply.github.com>
Date: Thu, 28 Dec 2023 00:24:25 -0600
Subject: [PATCH] Treat walruses like other binary operators in subscripts
 (#4109)

Fixes #4078
---
 CHANGES.md                          | 3 +++
 src/black/lines.py                  | 9 ++++++++-
 tests/data/cases/preview_pep_572.py | 4 ++--
 3 files changed, 13 insertions(+), 3 deletions(-)

diff --git a/CHANGES.md b/CHANGES.md
index 526cbd12123..1444463050f 100644
--- 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
index 0cd4189a778..d153b8c2e1b 100644
--- 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
index 8e801ff6cdc..75ad0cc4176 100644
--- 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)]