Implement support for PEP 646 (#3071)

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 -->
 

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
 
 

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 @@ class Feature(Enum):
         Feature.ANN_ASSIGN_EXTENDED_RHS,
         Feature.PATTERN_MATCHING,
         Feature.EXCEPT_STAR,
+        Feature.VARIADIC_GENERICS,
     },
 }
 

src/black/nodes.py

@@ -120,6 +120,7 @@
     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)

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))* [',']

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

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__

tests/test_black.py

@@ -804,6 +804,12 @@ def test_get_features_used(self) -> None:
         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 [