From 67fa8d4d542dcb2f1e85dbab045da5ac0f7c80da Mon Sep 17 00:00:00 2001 From: Charlie Marsh Date: Sat, 29 Jul 2023 16:45:55 -0400 Subject: [PATCH] Avoid removing whitespace for walrus operators within subscripts --- CHANGES.md | 1 + src/black/nodes.py | 3 +++ tests/data/py_310/pep_572_py310.py | 6 +++--- 3 files changed, 7 insertions(+), 3 deletions(-) diff --git a/CHANGES.md b/CHANGES.md index 709c767b329..9e535e37033 100644 --- a/CHANGES.md +++ b/CHANGES.md @@ -66,6 +66,7 @@ (#3740) - Fix error in AST validation when _Black_ removes trailing whitespace in a type comment (#3773) +- Fix a bug whereby spaces were removed from walrus operators within subscript (#3823) ### Preview style diff --git a/src/black/nodes.py b/src/black/nodes.py index 45423b2596b..87d13020a16 100644 --- a/src/black/nodes.py +++ b/src/black/nodes.py @@ -345,6 +345,9 @@ def whitespace(leaf: Leaf, *, complex_subscript: bool) -> str: # noqa: C901 return NO + elif t == token.COLONEQUAL or prev.type == token.COLONEQUAL: + return SPACE + elif not complex_subscript: return NO diff --git a/tests/data/py_310/pep_572_py310.py b/tests/data/py_310/pep_572_py310.py index cb82b2d23f8..3487ac8536a 100644 --- a/tests/data/py_310/pep_572_py310.py +++ b/tests/data/py_310/pep_572_py310.py @@ -1,7 +1,7 @@ # Unparenthesized walruses are now allowed in indices since Python 3.10. -x[a:=0] -x[a:=0, b:=1] -x[5, b:=0] +x[a := 0] +x[a := 0, b := 1] +x[5, b := 0] # Walruses are allowed inside generator expressions on function calls since 3.10. if any(match := pattern_error.match(s) for s in buffer):