From: Charlie Marsh Date: Fri, 8 Sep 2023 14:37:13 +0000 (+0200) Subject: Avoid removing whitespace for walrus operators within subscripts (#3823) X-Git-Url: https://git.madduck.net/etc/vim.git/commitdiff_plain/a20338cf100ff20a24e2058c6f6014e82efdf914?ds=inline Avoid removing whitespace for walrus operators within subscripts (#3823) Co-authored-by: hauntsaninja Co-authored-by: Jelle Zijlstra --- diff --git a/CHANGES.md b/CHANGES.md index af9fc49..4aa3123 100644 --- a/CHANGES.md +++ b/CHANGES.md @@ -80,6 +80,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/lines.py b/src/black/lines.py index 0a307b4..f3044ce 100644 --- a/src/black/lines.py +++ b/src/black/lines.py @@ -81,7 +81,9 @@ class Line: # Note: at this point leaf.prefix should be empty except for # imports, for which we only preserve newlines. leaf.prefix += whitespace( - leaf, complex_subscript=self.is_complex_subscript(leaf) + leaf, + complex_subscript=self.is_complex_subscript(leaf), + mode=self.mode, ) if self.inside_brackets or not preformatted or track_bracket: self.bracket_tracker.mark(leaf) diff --git a/src/black/mode.py b/src/black/mode.py index 282c166..06d20b7 100644 --- a/src/black/mode.py +++ b/src/black/mode.py @@ -183,6 +183,7 @@ class Preview(Enum): wrap_long_dict_values_in_parens = auto() wrap_multiple_context_managers_in_parens = auto() dummy_implementations = auto() + walrus_subscript = auto() class Deprecated(UserWarning): diff --git a/src/black/nodes.py b/src/black/nodes.py index 45423b2..edd201a 100644 --- a/src/black/nodes.py +++ b/src/black/nodes.py @@ -13,6 +13,7 @@ else: from mypy_extensions import mypyc_attr from black.cache import CACHE_DIR +from black.mode import Mode, Preview from black.strings import has_triple_quotes from blib2to3 import pygram from blib2to3.pgen2 import token @@ -171,7 +172,7 @@ class Visitor(Generic[T]): yield from self.visit(child) -def whitespace(leaf: Leaf, *, complex_subscript: bool) -> str: # noqa: C901 +def whitespace(leaf: Leaf, *, complex_subscript: bool, mode: Mode) -> str: # noqa: C901 """Return whitespace prefix if needed for the given `leaf`. `complex_subscript` signals whether the given leaf is part of a subscription @@ -345,6 +346,11 @@ def whitespace(leaf: Leaf, *, complex_subscript: bool) -> str: # noqa: C901 return NO + elif Preview.walrus_subscript in mode and ( + t == token.COLONEQUAL or prev.type == token.COLONEQUAL + ): + return SPACE + elif not complex_subscript: return NO diff --git a/tests/data/preview/pep_572.py b/tests/data/preview/pep_572.py new file mode 100644 index 0000000..a50e130 --- /dev/null +++ b/tests/data/preview/pep_572.py @@ -0,0 +1,6 @@ +x[(a:=0):] +x[:(a:=0)] + +# output +x[(a := 0):] +x[:(a := 0)] diff --git a/tests/data/preview_py_310/pep_572.py b/tests/data/preview_py_310/pep_572.py new file mode 100644 index 0000000..78d4e9e --- /dev/null +++ b/tests/data/preview_py_310/pep_572.py @@ -0,0 +1,12 @@ +x[a:=0] +x[a := 0] +x[a := 0, b := 1] +x[5, b := 0] +x[a:=0,b:=1] + +# output +x[a := 0] +x[a := 0] +x[a := 0, b := 1] +x[5, b := 0] +x[a := 0, b := 1] diff --git a/tests/test_format.py b/tests/test_format.py index fb4d8eb..0650a2d 100644 --- a/tests/test_format.py +++ b/tests/test_format.py @@ -56,6 +56,13 @@ def test_preview_context_managers_targeting_py39() -> None: assert_format(source, expected, mode, minimum_version=(3, 9)) +@pytest.mark.parametrize("filename", all_data_cases("preview_py_310")) +def test_preview_python_310(filename: str) -> None: + source, expected = read_data("preview_py_310", filename) + mode = black.Mode(target_versions={black.TargetVersion.PY310}, preview=True) + assert_format(source, expected, mode, minimum_version=(3, 10)) + + @pytest.mark.parametrize( "filename", all_data_cases("preview_context_managers/auto_detect") )