From 366a0806eb4edd94d4d20cc6912333586bdc9fa7 Mon Sep 17 00:00:00 2001
From: Richard Si <63936253+ichard26@users.noreply.github.com>
Date: Thu, 26 Aug 2021 16:59:01 -0400
Subject: [PATCH] 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                  | 2 ++
 src/blib2to3/Grammar.txt    | 6 +++---
 src/blib2to3/README         | 7 ++++++-
 tests/data/pep_572_py310.py | 4 ++++
 tests/data/pep_572_py39.py  | 7 +++++++
 tests/test_black.py         | 9 +++++++++
 6 files changed, 31 insertions(+), 4 deletions(-)
 create mode 100644 tests/data/pep_572_py310.py
 create mode 100644 tests/data/pep_572_py39.py

diff --git a/CHANGES.md b/CHANGES.md
index 47e64cf..d28d766 100644
--- a/CHANGES.md
+++ b/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_
 
diff --git a/src/blib2to3/Grammar.txt b/src/blib2to3/Grammar.txt
index 69b9af9..ac8a067 100644
--- a/src/blib2to3/Grammar.txt
+++ b/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
 
diff --git a/src/blib2to3/README b/src/blib2to3/README
index a43f15c..ccad283 100644
--- a/src/blib2to3/README
+++ b/src/blib2to3/README
@@ -13,4 +13,9 @@ Reasons for forking:
 - ability to Cythonize
 
 Change Log:
-- Changes default logger used by Driver
\ No newline at end of file
+- 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
diff --git a/tests/data/pep_572_py310.py b/tests/data/pep_572_py310.py
new file mode 100644
index 0000000..2aef589
--- /dev/null
+++ b/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]
diff --git a/tests/data/pep_572_py39.py b/tests/data/pep_572_py39.py
new file mode 100644
index 0000000..7bbd509
--- /dev/null
+++ b/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)]
diff --git a/tests/test_black.py b/tests/test_black.py
index 5c72050..8a37f7c 100644
--- a/tests/test_black.py
+++ b/tests/test_black.py
@@ -26,6 +26,7 @@ from typing import (
 import pytest
 import unittest
 from unittest.mock import patch, MagicMock
+from parameterized import parameterized
 
 import click
 from click import unstyle
@@ -299,6 +300,14 @@ class BlackTestCase(BlackBaseTestCase):
         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))
-- 
2.39.5