Parse type parameters in class definitions

Author: zanieb

parser/src/parser.rs

@@ -636,6 +636,79 @@ class Foo(A, B):
         insta::assert_debug_snapshot!(ast::Suite::parse(source, "<test>").unwrap());
     }
 
+
+    #[test]
+    #[cfg(feature = "all-nodes-with-ranges")]
+    fn test_parse_class_with_empty_generic() {
+        let source = "\
+class Foo[](A, B):
+  pass
+";
+        insta::assert_debug_snapshot!(ast::Suite::parse(source, "<test>").unwrap());
+    }
+
+
+    #[test]
+    #[cfg(feature = "all-nodes-with-ranges")]
+    fn test_parse_class_with_generic_type() {
+        let source = "\
+class Foo[T](A, B):
+  pass
+";
+        insta::assert_debug_snapshot!(ast::Suite::parse(source, "<test>").unwrap());
+    }
+
+    #[test]
+    #[cfg(feature = "all-nodes-with-ranges")]
+    fn test_parse_class_with_generic_type_with_bound() {
+        let source = "\
+class Foo[T: str](A, B):
+  pass
+";
+        insta::assert_debug_snapshot!(ast::Suite::parse(source, "<test>").unwrap());
+    }
+
+    #[test]
+    #[cfg(feature = "all-nodes-with-ranges")]
+    fn test_parse_class_with_multiple_generic_types() {
+        let source = "\
+class Foo[T, U](A, B):
+  pass
+";
+        insta::assert_debug_snapshot!(ast::Suite::parse(source, "<test>").unwrap());
+    }
+
+    #[test]
+    #[cfg(feature = "all-nodes-with-ranges")]
+    fn test_parse_class_with_generic_type_var_tuple() {
+        let source = "\
+class Foo[*U](A, B):
+  pass
+";
+        insta::assert_debug_snapshot!(ast::Suite::parse(source, "<test>").unwrap());
+    }
+
+
+    #[test]
+    #[cfg(feature = "all-nodes-with-ranges")]
+    fn test_parse_class_with_generic_param_spec() {
+        let source = "\
+class Foo[**P](A, B):
+  pass
+";
+        insta::assert_debug_snapshot!(ast::Suite::parse(source, "<test>").unwrap());
+    }
+
+    #[test]
+    #[cfg(feature = "all-nodes-with-ranges")]
+    fn test_parse_class_with_all_possible_generic_types() {
+        let source = "\
+class Foo[X, Y, *U, **P](A, B):
+  pass
+";
+        insta::assert_debug_snapshot!(ast::Suite::parse(source, "<test>").unwrap());
+    }
+
     #[test]
     #[cfg(feature = "all-nodes-with-ranges")]
     fn test_parse_dict_comprehension() {

parser/src/python.lalrpop

@@ -1126,27 +1126,51 @@ KwargParameter<ArgType>: Option<Box<ast::Arg>> = {
 };
 
 ClassDef: ast::Stmt = {
-    <decorator_list:Decorator*> <location:@L> "class" <name:Identifier> <a:("(" ArgumentList ")")?> ":" <body:Suite> => {
+    <decorator_list:Decorator*> <location:@L> "class" <name:Identifier> <type_params:TypeParamList?> <a:("(" ArgumentList ")")?> ":" <body:Suite> => {
         let (bases, keywords) = match a {
             Some((_, arg, _)) => (arg.args, arg.keywords),
             None => (vec![], vec![]),
         };
         let end_location = body.last().unwrap().end();
-        let type_params = Vec::new();
         ast::Stmt::ClassDef(
             ast::StmtClassDef {
                 name,
                 bases,
                 keywords,
                 body,
                 decorator_list,
-                type_params,
+                type_params: type_params.unwrap_or_default(),
                 range: (location..end_location).into()
             },
         )
     },
 };
 
+
+TypeParamList: Vec<ast::TypeParam> = {
+    <location:@L> "[" <vars:OneOrMore<TypeParam>> ","? "]" <end_location:@R> => {
+        vars
+    }
+};
+
+TypeParam: ast::TypeParam = {
+    <location:@L> <name:Identifier> <bound:(":" <Test<"all">>)?> <end_location:@R>  => {
+        ast::TypeParam::TypeVar(
+            ast::TypeParamTypeVar { name, bound: bound.map(Box::new), range: (location..end_location).into() }
+        )
+    },
+    <location:@L> "*" <name:Identifier> <end_location:@R>  => {
+        ast::TypeParam::TypeVarTuple(
+            ast::TypeParamTypeVarTuple { name, range: (location..end_location).into() }
+        )
+    },
+    <location:@L> "**" <name:Identifier> <end_location:@R>  => {
+        ast::TypeParam::ParamSpec(
+            ast::TypeParamParamSpec { name, range: (location..end_location).into() }
+        )
+    }
+};
+
 // Decorators:
 Decorator: ast::Expr = {
     <location:@L> "@" <p:NamedExpressionTest> "\n" => {