Implement support for PEP 646 (#3071)

diff --git a/CHANGES.md b/CHANGES.md
index 8f43431c8421c23d9986697a1f1295fd848f06d4..6bc67f9db06babe585d200dd859ebee84226c5e6 100644 (file)
--- a/CHANGES.md
+++ b/CHANGES.md
@@ -56,6 +56,8 @@
 
 - [PEP 654](https://peps.python.org/pep-0654/#except) syntax (for example,
   `except *ExceptionGroup:`) is now supported (#3016)
+- [PEP 646](https://peps.python.org/pep-0646) syntax (for example,
+  `Array[Batch, *Shape]` or `def fn(*args: *T) -> None`) is now supported (#3071)
 
 <!-- Changes to the parser or to version autodetection -->
 

diff --git a/src/black/__init__.py b/src/black/__init__.py
index 75321c3f35c96c2ebfc2e62d755ac0500ad16a73..8872102a6ea8baec5d74094297bb327adaccd8c7 100644 (file)
--- a/src/black/__init__.py
+++ b/src/black/__init__.py
@@ -1308,6 +1308,18 @@ def get_features_used(  # noqa: C901
         ):
             features.add(Feature.EXCEPT_STAR)
 
+        elif n.type in {syms.subscriptlist, syms.trailer} and any(
+            child.type == syms.star_expr for child in n.children
+        ):
+            features.add(Feature.VARIADIC_GENERICS)
+
+        elif (
+            n.type == syms.tname_star
+            and len(n.children) == 3
+            and n.children[2].type == syms.star_expr
+        ):
+            features.add(Feature.VARIADIC_GENERICS)
+
     return features
 
 

diff --git a/src/black/mode.py b/src/black/mode.py
index a418e0eb66510ef6ea451cbc51f8b4be035816d4..bf79f6a31489276bcd79a611650db0d20038332a 100644 (file)
--- a/src/black/mode.py
+++ b/src/black/mode.py
@@ -49,6 +49,7 @@ class Feature(Enum):
     UNPACKING_ON_FLOW = 12
     ANN_ASSIGN_EXTENDED_RHS = 13
     EXCEPT_STAR = 14
+    VARIADIC_GENERICS = 15
     FORCE_OPTIONAL_PARENTHESES = 50
 
     # __future__ flags
@@ -132,6 +133,7 @@ VERSION_TO_FEATURES: Dict[TargetVersion, Set[Feature]] = {
         Feature.ANN_ASSIGN_EXTENDED_RHS,
         Feature.PATTERN_MATCHING,
         Feature.EXCEPT_STAR,
+        Feature.VARIADIC_GENERICS,
     },
 }
 

diff --git a/src/black/nodes.py b/src/black/nodes.py
index 37b96a498d63f8ae067a555b8fb51642381a6181..918038f69ba4491c7e47492af0f0b3c6f80a3600 100644 (file)
--- a/src/black/nodes.py
+++ b/src/black/nodes.py
@@ -120,6 +120,7 @@ TEST_DESCENDANTS: Final = {
     syms.term,
     syms.power,
 }
+TYPED_NAMES: Final = {syms.tname, syms.tname_star}
 ASSIGNMENTS: Final = {
     "=",
     "+=",
@@ -243,6 +244,14 @@ def whitespace(leaf: Leaf, *, complex_subscript: bool) -> str:  # noqa: C901
                     # that, too.
                     return prevp.prefix
 
+        elif (
+            prevp.type == token.STAR
+            and parent_type(prevp) == syms.star_expr
+            and parent_type(prevp.parent) == syms.subscriptlist
+        ):
+            # No space between typevar tuples.
+            return NO
+
         elif prevp.type in VARARGS_SPECIALS:
             if is_vararg(prevp, within=VARARGS_PARENTS | UNPACKING_PARENTS):
                 return NO
@@ -281,7 +290,7 @@ def whitespace(leaf: Leaf, *, complex_subscript: bool) -> str:  # noqa: C901
             return NO
 
         if t == token.EQUAL:
-            if prev.type != syms.tname:
+            if prev.type not in TYPED_NAMES:
                 return NO
 
         elif prev.type == token.EQUAL:
@@ -292,7 +301,7 @@ def whitespace(leaf: Leaf, *, complex_subscript: bool) -> str:  # noqa: C901
         elif prev.type != token.COMMA:
             return NO
 
-    elif p.type == syms.tname:
+    elif p.type in TYPED_NAMES:
         # type names
         if not prev:
             prevp = preceding_leaf(p)

diff --git a/src/blib2to3/Grammar.txt b/src/blib2to3/Grammar.txt
index 1de5416551399f54fd59016f2bc1ba7deda0bb9f..ac7ad7643ffb8deaef567e85849d7da140143a98 100644 (file)
--- a/src/blib2to3/Grammar.txt
+++ b/src/blib2to3/Grammar.txt
@@ -24,7 +24,7 @@ parameters: '(' [typedargslist] ')'
 #     arguments = argument (',' argument)*
 #     argument = tfpdef ['=' test]
 #     kwargs = '**' tname [',']
-#     args = '*' [tname]
+#     args = '*' [tname_star]
 #     kwonly_kwargs = (',' argument)* [',' [kwargs]]
 #     args_kwonly_kwargs = args kwonly_kwargs | kwargs
 #     poskeyword_args_kwonly_kwargs = arguments [',' [args_kwonly_kwargs]]
@@ -34,14 +34,15 @@ parameters: '(' [typedargslist] ')'
 # It needs to be fully expanded to allow our LL(1) parser to work on it.
 
 typedargslist: tfpdef ['=' test] (',' tfpdef ['=' test])* ',' '/' [
-                     ',' [((tfpdef ['=' test] ',')* ('*' [tname] (',' tname ['=' test])*
+                     ',' [((tfpdef ['=' test] ',')* ('*' [tname_star] (',' tname ['=' test])*
                             [',' ['**' tname [',']]] | '**' tname [','])
                      | tfpdef ['=' test] (',' tfpdef ['=' test])* [','])]
-                ] | ((tfpdef ['=' test] ',')* ('*' [tname] (',' tname ['=' test])*
+                ] | ((tfpdef ['=' test] ',')* ('*' [tname_star] (',' tname ['=' test])*
                      [',' ['**' tname [',']]] | '**' tname [','])
                      | tfpdef ['=' test] (',' tfpdef ['=' test])* [','])
 
 tname: NAME [':' test]
+tname_star: NAME [':' (test|star_expr)]
 tfpdef: tname | '(' tfplist ')'
 tfplist: tfpdef (',' tfpdef)* [',']
 
@@ -163,7 +164,7 @@ listmaker: (namedexpr_test|star_expr) ( old_comp_for | (',' (namedexpr_test|star
 testlist_gexp: (namedexpr_test|star_expr) ( old_comp_for | (',' (namedexpr_test|star_expr))* [','] )
 lambdef: 'lambda' [varargslist] ':' test
 trailer: '(' [arglist] ')' | '[' subscriptlist ']' | '.' NAME
-subscriptlist: subscript (',' subscript)* [',']
+subscriptlist: (subscript|star_expr) (',' (subscript|star_expr))* [',']
 subscript: test [':=' test] | [test] ':' [test] [sliceop]
 sliceop: ':' [test]
 exprlist: (expr|star_expr) (',' (expr|star_expr))* [',']

diff --git a/src/blib2to3/pygram.py b/src/blib2to3/pygram.py
index a3df9be1265655a413cdfae0163ab14512f85df1..99012cdd9cb30f05369b6046702b73829bed2016 100644 (file)
--- a/src/blib2to3/pygram.py
+++ b/src/blib2to3/pygram.py
@@ -123,6 +123,7 @@ class _python_symbols(Symbols):
     tfpdef: int
     tfplist: int
     tname: int
+    tname_star: int
     trailer: int
     try_stmt: int
     typedargslist: int

diff --git a/tests/data/py_311/pep_646.py b/tests/data/py_311/pep_646.py
new file mode 100644 (file)
index 0000000..e843ecf
--- /dev/null
+++ b/tests/data/py_311/pep_646.py
@@ -0,0 +1,194 @@
+A[*b]
+A[*b] = 1
+A
+del A[*b]
+A
+A[*b, *b]
+A[*b, *b] = 1
+A
+del A[*b, *b]
+A
+A[b, *b]
+A[b, *b] = 1
+A
+del A[b, *b]
+A
+A[*b, b]
+A[*b, b] = 1
+A
+del A[*b, b]
+A
+A[b, b, *b]
+A[b, b, *b] = 1
+A
+del A[b, b, *b]
+A
+A[*b, b, b]
+A[*b, b, b] = 1
+A
+del A[*b, b, b]
+A
+A[b, *b, b]
+A[b, *b, b] = 1
+A
+del A[b, *b, b]
+A
+A[b, b, *b, b]
+A[b, b, *b, b] = 1
+A
+del A[b, b, *b, b]
+A
+A[b, *b, b, b]
+A[b, *b, b, b] = 1
+A
+del A[b, *b, b, b]
+A
+A[A[b, *b, b]]
+A[A[b, *b, b]] = 1
+A
+del A[A[b, *b, b]]
+A
+A[*A[b, *b, b]]
+A[*A[b, *b, b]] = 1
+A
+del A[*A[b, *b, b]]
+A
+A[b, ...]
+A[b, ...] = 1
+A
+del A[b, ...]
+A
+A[*A[b, ...]]
+A[*A[b, ...]] = 1
+A
+del A[*A[b, ...]]
+A
+l = [1, 2, 3]
+A[*l]
+A[*l] = 1
+A
+del A[*l]
+A
+A[*l, 4]
+A[*l, 4] = 1
+A
+del A[*l, 4]
+A
+A[0, *l]
+A[0, *l] = 1
+A
+del A[0, *l]
+A
+A[1:2, *l]
+A[1:2, *l] = 1
+A
+del A[1:2, *l]
+A
+repr(A[1:2, *l]) == repr(A[1:2, 1, 2, 3])
+t = (1, 2, 3)
+A[*t]
+A[*t] = 1
+A
+del A[*t]
+A
+A[*t, 4]
+A[*t, 4] = 1
+A
+del A[*t, 4]
+A
+A[0, *t]
+A[0, *t] = 1
+A
+del A[0, *t]
+A
+A[1:2, *t]
+A[1:2, *t] = 1
+A
+del A[1:2, *t]
+A
+repr(A[1:2, *t]) == repr(A[1:2, 1, 2, 3])
+
+
+def returns_list():
+    return [1, 2, 3]
+
+
+A[returns_list()]
+A[returns_list()] = 1
+A
+del A[returns_list()]
+A
+A[returns_list(), 4]
+A[returns_list(), 4] = 1
+A
+del A[returns_list(), 4]
+A
+A[*returns_list()]
+A[*returns_list()] = 1
+A
+del A[*returns_list()]
+A
+A[*returns_list(), 4]
+A[*returns_list(), 4] = 1
+A
+del A[*returns_list(), 4]
+A
+A[0, *returns_list()]
+A[0, *returns_list()] = 1
+A
+del A[0, *returns_list()]
+A
+A[*returns_list(), *returns_list()]
+A[*returns_list(), *returns_list()] = 1
+A
+del A[*returns_list(), *returns_list()]
+A
+A[1:2, *b]
+A[*b, 1:2]
+A[1:2, *b, 1:2]
+A[*b, 1:2, *b]
+A[1:, *b]
+A[*b, 1:]
+A[1:, *b, 1:]
+A[*b, 1:, *b]
+A[:1, *b]
+A[*b, :1]
+A[:1, *b, :1]
+A[*b, :1, *b]
+A[:, *b]
+A[*b, :]
+A[:, *b, :]
+A[*b, :, *b]
+A[a * b()]
+A[a * b(), *c, *d(), e * f(g * h)]
+A[a * b(), :]
+A[a * b(), *c, *d(), e * f(g * h) :]
+A[[b] * len(c), :]
+
+
+def f1(*args: *b):
+    pass
+
+
+f1.__annotations__
+
+
+def f2(*args: *b, arg1):
+    pass
+
+
+f2.__annotations__
+
+
+def f3(*args: *b, arg1: int):
+    pass
+
+
+f3.__annotations__
+
+
+def f4(*args: *b, arg1: int = 2):
+    pass
+
+
+f4.__annotations__

diff --git a/tests/test_black.py b/tests/test_black.py
index a633e678dd700d70f17ad622bd21cb0cc602ade0..02a707e8996aa25beab3137a91281623126f0af5 100644 (file)
--- a/tests/test_black.py
+++ b/tests/test_black.py
@@ -804,6 +804,12 @@ class BlackTestCase(BlackBaseTestCase):
         self.assertEqual(black.get_features_used(node), set())
         node = black.lib2to3_parse("try: pass\nexcept *Group: pass")
         self.assertEqual(black.get_features_used(node), {Feature.EXCEPT_STAR})
+        node = black.lib2to3_parse("a[*b]")
+        self.assertEqual(black.get_features_used(node), {Feature.VARIADIC_GENERICS})
+        node = black.lib2to3_parse("a[x, *y(), z] = t")
+        self.assertEqual(black.get_features_used(node), {Feature.VARIADIC_GENERICS})
+        node = black.lib2to3_parse("def fn(*args: *T): pass")
+        self.assertEqual(black.get_features_used(node), {Feature.VARIADIC_GENERICS})
 
     def test_get_features_used_for_future_flags(self) -> None:
         for src, features in [