From 0d7b94817d47e43b6ef3e0d11b134b8ca6867acf Mon Sep 17 00:00:00 2001
From: Charlie Marsh <charlie.r.marsh@gmail.com>
Date: Thu, 23 Feb 2023 08:37:44 -0500
Subject: [PATCH] Allow type variable tuple for *args

---
 parser/python.lalrpop                         |  21 +-
 parser/src/parser.rs                          |  12 ++
 ...ser__parser__tests__variadic_generics.snap | 188 ++++++++++++++++++
 3 files changed, 214 insertions(+), 7 deletions(-)
 create mode 100644 parser/src/snapshots/rustpython_parser__parser__tests__variadic_generics.snap

diff --git a/parser/python.lalrpop b/parser/python.lalrpop
index ee7f297..4eceb5b 100644
--- a/parser/python.lalrpop
+++ b/parser/python.lalrpop
@@ -1180,7 +1180,7 @@ FuncDef: ast::Stmt = {
 };
 
 Parameters: ast::Arguments = {
-    "(" <a: (ParameterList<TypedParameter>)?> ")" =>? {
+    "(" <a: (ParameterList<TypedParameter, StarTypedParameter>)?> ")" =>? {
         let args = validate_arguments(
             a.unwrap_or_else(|| ast::Arguments {
                 posonlyargs: vec![],
@@ -1199,8 +1199,8 @@ Parameters: ast::Arguments = {
 
 // Note that this is a macro which is used once for function defs, and
 // once for lambda defs.
-ParameterList<ArgType>: ast::Arguments = {
-    <param1:ParameterDefs<ArgType>> <args2:("," ParameterListStarArgs<ArgType>)?> ","? =>? {
+ParameterList<ArgType, StarArgType>: ast::Arguments = {
+    <param1:ParameterDefs<ArgType>> <args2:("," ParameterListStarArgs<ArgType, StarArgType>)?> ","? =>? {
         let (posonlyargs, args, defaults) = parse_params(param1)?;
 
         // Now gather rest of parameters:
@@ -1235,7 +1235,7 @@ ParameterList<ArgType>: ast::Arguments = {
             kw_defaults,
         })
     },
-    <params:ParameterListStarArgs<ArgType>> ","? => {
+    <params:ParameterListStarArgs<ArgType, StarArgType>> ","? => {
         let (vararg, kwonlyargs, kw_defaults, kwarg) = params;
         ast::Arguments {
             posonlyargs: vec![],
@@ -1291,11 +1291,18 @@ TypedParameter: ast::Arg = {
     },
 };
 
+StarTypedParameter: ast::Arg = {
+    <location:@L> <arg:Identifier> <a:(":" TestOrStarExpr)?> <end_location:@R> => {
+        let annotation = a.map(|x| Box::new(x.1));
+        ast::Arg::new(location, end_location, ast::ArgData { arg, annotation, type_comment: None })
+    },
+};
+
 // Use inline here to make sure the "," is not creating an ambiguity.
 // TODO: figure out another grammar that makes this inline no longer required.
 #[inline]
-ParameterListStarArgs<ArgType>: (Option<Box<ast::Arg>>, Vec<ast::Arg>, Vec<ast::Expr>, Option<Box<ast::Arg>>) = {
-    <location:@L> "*" <va:ArgType?> <kw:("," ParameterDef<ArgType>)*> <kwarg:("," KwargParameter<ArgType>)?> =>? {
+ParameterListStarArgs<ArgType, StarArgType>: (Option<Box<ast::Arg>>, Vec<ast::Arg>, Vec<ast::Expr>, Option<Box<ast::Arg>>) = {
+    <location:@L> "*" <va:StarArgType?> <kw:("," ParameterDef<ArgType>)*> <kwarg:("," KwargParameter<ArgType>)?> =>? {
         // Extract keyword arguments:
         let mut kwonlyargs = Vec::new();
         let mut kw_defaults = Vec::new();
@@ -1413,7 +1420,7 @@ NamedExpression: ast::Expr = {
 };
 
 LambdaDef: ast::Expr = {
-    <location:@L> "lambda" <p:ParameterList<UntypedParameter>?> ":" <body:Test<"all">> <end_location:@R> =>? {
+    <location:@L> "lambda" <p:ParameterList<UntypedParameter, UntypedParameter>?> ":" <body:Test<"all">> <end_location:@R> =>? {
         let p = validate_arguments(
             p.unwrap_or_else(|| {
                 ast::Arguments {
diff --git a/parser/src/parser.rs b/parser/src/parser.rs
index 50c705f..c2c4a38 100644
--- a/parser/src/parser.rs
+++ b/parser/src/parser.rs
@@ -803,6 +803,18 @@ match x:
 match x:
     case (0,):
         y = 0
+"#,
+            "<test>",
+        )
+        .unwrap();
+        insta::assert_debug_snapshot!(parse_ast);
+    }
+
+    #[test]
+    fn test_variadic_generics() {
+        let parse_ast = parse_program(
+            r#"
+def args_to_tuple(*args: *Ts) -> Tuple[*Ts]: ...
 "#,
             "<test>",
         )
diff --git a/parser/src/snapshots/rustpython_parser__parser__tests__variadic_generics.snap b/parser/src/snapshots/rustpython_parser__parser__tests__variadic_generics.snap
new file mode 100644
index 0000000..a6875c3
--- /dev/null
+++ b/parser/src/snapshots/rustpython_parser__parser__tests__variadic_generics.snap
@@ -0,0 +1,188 @@
+---
+source: compiler/parser/src/parser.rs
+expression: parse_ast
+---
+[
+    Located {
+        location: Location {
+            row: 2,
+            column: 0,
+        },
+        end_location: Some(
+            Location {
+                row: 2,
+                column: 48,
+            },
+        ),
+        custom: (),
+        node: FunctionDef {
+            name: "args_to_tuple",
+            args: Arguments {
+                posonlyargs: [],
+                args: [],
+                vararg: Some(
+                    Located {
+                        location: Location {
+                            row: 2,
+                            column: 19,
+                        },
+                        end_location: Some(
+                            Location {
+                                row: 2,
+                                column: 28,
+                            },
+                        ),
+                        custom: (),
+                        node: ArgData {
+                            arg: "args",
+                            annotation: Some(
+                                Located {
+                                    location: Location {
+                                        row: 2,
+                                        column: 25,
+                                    },
+                                    end_location: Some(
+                                        Location {
+                                            row: 2,
+                                            column: 28,
+                                        },
+                                    ),
+                                    custom: (),
+                                    node: Starred {
+                                        value: Located {
+                                            location: Location {
+                                                row: 2,
+                                                column: 26,
+                                            },
+                                            end_location: Some(
+                                                Location {
+                                                    row: 2,
+                                                    column: 28,
+                                                },
+                                            ),
+                                            custom: (),
+                                            node: Name {
+                                                id: "Ts",
+                                                ctx: Load,
+                                            },
+                                        },
+                                        ctx: Load,
+                                    },
+                                },
+                            ),
+                            type_comment: None,
+                        },
+                    },
+                ),
+                kwonlyargs: [],
+                kw_defaults: [],
+                kwarg: None,
+                defaults: [],
+            },
+            body: [
+                Located {
+                    location: Location {
+                        row: 2,
+                        column: 45,
+                    },
+                    end_location: Some(
+                        Location {
+                            row: 2,
+                            column: 48,
+                        },
+                    ),
+                    custom: (),
+                    node: Expr {
+                        value: Located {
+                            location: Location {
+                                row: 2,
+                                column: 45,
+                            },
+                            end_location: Some(
+                                Location {
+                                    row: 2,
+                                    column: 48,
+                                },
+                            ),
+                            custom: (),
+                            node: Constant {
+                                value: Ellipsis,
+                                kind: None,
+                            },
+                        },
+                    },
+                },
+            ],
+            decorator_list: [],
+            returns: Some(
+                Located {
+                    location: Location {
+                        row: 2,
+                        column: 33,
+                    },
+                    end_location: Some(
+                        Location {
+                            row: 2,
+                            column: 43,
+                        },
+                    ),
+                    custom: (),
+                    node: Subscript {
+                        value: Located {
+                            location: Location {
+                                row: 2,
+                                column: 33,
+                            },
+                            end_location: Some(
+                                Location {
+                                    row: 2,
+                                    column: 38,
+                                },
+                            ),
+                            custom: (),
+                            node: Name {
+                                id: "Tuple",
+                                ctx: Load,
+                            },
+                        },
+                        slice: Located {
+                            location: Location {
+                                row: 2,
+                                column: 39,
+                            },
+                            end_location: Some(
+                                Location {
+                                    row: 2,
+                                    column: 42,
+                                },
+                            ),
+                            custom: (),
+                            node: Starred {
+                                value: Located {
+                                    location: Location {
+                                        row: 2,
+                                        column: 40,
+                                    },
+                                    end_location: Some(
+                                        Location {
+                                            row: 2,
+                                            column: 42,
+                                        },
+                                    ),
+                                    custom: (),
+                                    node: Name {
+                                        id: "Ts",
+                                        ctx: Load,
+                                    },
+                                },
+                                ctx: Load,
+                            },
+                        },
+                        ctx: Load,
+                    },
+                },
+            ),
+            type_comment: None,
+        },
+    },
+]