blib2to3: support unparenthesized wulruses in more places (#2447)

Implementation stolen from PR davidhalter/parso#162. Thanks parso!

I could add support for these newer syntactical constructs in the
target version detection logic, but until I get diff-shades up
and running I don't feel very comfortable adding the code.

CHANGES.md

@@ -11,6 +11,8 @@
   hardened to handle more edge cases during quote normalization (#2437)
 - Avoid changing a function return type annotation's type to a tuple by adding a
   trailing comma (#2384)
+- Parsing support has been added for unparenthesized walruses in set literals, set
+  comprehensions, and indices (#2447).
 
 ### _Blackd_
 

src/blib2to3/Grammar.txt

@@ -157,14 +157,14 @@ testlist_gexp: (namedexpr_test|star_expr) ( old_comp_for | (',' (namedexpr_test|
 lambdef: 'lambda' [varargslist] ':' test
 trailer: '(' [arglist] ')' | '[' subscriptlist ']' | '.' NAME
 subscriptlist: subscript (',' subscript)* [',']
-subscript: test | [test] ':' [test] [sliceop]
+subscript: test [':=' test] | [test] ':' [test] [sliceop]
 sliceop: ':' [test]
 exprlist: (expr|star_expr) (',' (expr|star_expr))* [',']
 testlist: test (',' test)* [',']
 dictsetmaker: ( ((test ':' test | '**' expr)
                  (comp_for | (',' (test ':' test | '**' expr))* [','])) |
-                ((test | star_expr)
-		 (comp_for | (',' (test | star_expr))* [','])) )
+                ((test [':=' test] | star_expr)
+		 (comp_for | (',' (test [':=' test] | star_expr))* [','])) )
 
 classdef: 'class' NAME ['(' [arglist] ')'] ':' suite
 

src/blib2to3/README

@@ -13,4 +13,9 @@ Reasons for forking:
 - ability to Cythonize
 
 Change Log:
-- Changes default logger used by Driver
+- Changes default logger used by Driver
+- Backported the following upstream parser changes:
+  - "bpo-42381: Allow walrus in set literals and set comprehensions (GH-23332)"
+    https://github.com/python/cpython/commit/cae60187cf7a7b26281d012e1952fafe4e2e97e9
+  - "bpo-42316: Allow unparenthesized walrus operator in indexes (GH-23317)"
+    https://github.com/python/cpython/commit/b0aba1fcdc3da952698d99aec2334faa79a8b68c

tests/data/pep_572_py310.py

@@ -0,0 +1,4 @@
+# Unparenthesized walruses are now allowed in indices since Python 3.10.
+x[a:=0]
+x[a:=0, b:=1]
+x[5, b:=0]

tests/data/pep_572_py39.py

@@ -0,0 +1,7 @@
+# Unparenthesized walruses are now allowed in set literals & set comprehensions
+# since Python 3.9
+{x := 1, 2, 3}
+{x4 := x ** 5 for x in range(7)}
+# We better not remove the parentheses here (since it's a 3.10 feature)
+x[(a := 1)]
+x[(a := 1), (b := 3)]

tests/test_black.py

@@ -26,6 +26,7 @@
 import pytest
 import unittest
 from unittest.mock import patch, MagicMock
+from parameterized import parameterized
 
 import click
 from click import unstyle
@@ -299,6 +300,14 @@ def test_pep_572_version_detection(self) -> None:
         versions = black.detect_target_versions(root)
         self.assertIn(black.TargetVersion.PY38, versions)
 
+    @parameterized.expand([(3, 9), (3, 10)])
+    def test_pep_572_newer_syntax(self, major: int, minor: int) -> None:
+        source, expected = read_data(f"pep_572_py{major}{minor}")
+        actual = fs(source, mode=DEFAULT_MODE)
+        self.assertFormatEqual(expected, actual)
+        if sys.version_info >= (major, minor):
+            black.assert_equivalent(source, actual)
+
     def test_expression_ff(self) -> None:
         source, expected = read_data("expression")
         tmp_file = Path(black.dump_to_file(source))