Add parsing of type alias statements

Author: zanieb

parser/build.rs

@@ -154,6 +154,7 @@ fn gen_phf(out_dir: &Path) {
         .entry("raise", "Tok::Raise")
         .entry("return", "Tok::Return")
         .entry("try", "Tok::Try")
+        .entry("type", "Tok::Type")
         .entry("while", "Tok::While")
         .entry("with", "Tok::With")
         .entry("yield", "Tok::Yield")

parser/src/parser.rs

@@ -636,6 +636,7 @@ 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_generic_types() {
@@ -867,11 +868,45 @@ except* OSError as e:
         assert!(parse(source, Mode::Interactive, "<embedded>").is_ok());
     }
 
+
+    #[test]
+    #[cfg(feature = "all-nodes-with-ranges")]
+    fn test_parse_type_declaration() {
+        let source = "\
+# A non-generic type alias
+type IntOrStr = int | str
+
+# A generic type alias
+type ListOrSet[T] = list[T] | set[T]
+
+# A type alias that includes a forward reference
+type AnimalOrVegetable = Animal | \"Vegetable\"
+
+# A generic self-referential type alias
+type RecursiveList[T] = T | list[RecursiveList[T]]
+";
+        insta::assert_debug_snapshot!(ast::Suite::parse(source, "<test>").unwrap());
+    }
+
+    #[test]
+    #[cfg(feature = "all-nodes-with-ranges")]
+    fn test_type_as_identifier() {
+        let source = r#"\
+type = lambda query: query == event
+print(type(12))
+"#;
+
+    use crate::lexer::lex;
+    let lexer = lex(source, Mode::Module);
+    println!("tokens {:#?}", lexer.map(|x| x.unwrap().0).collect::<Vec<_>>());
+
+    insta::assert_debug_snapshot!(ast::Suite::parse(source, "<test>").unwrap());
+    }
+
     #[test]
     #[cfg(feature = "all-nodes-with-ranges")]
     fn test_match_as_identifier() {
-        let parse_ast = ast::Suite::parse(
-            r#"
+        let source = r#"\
 match *a + b, c   # ((match * a) + b), c
 match *(a + b), c   # (match * (a + b)), c
 match (*a + b, c)   # match ((*(a + b)), c)
@@ -893,11 +928,8 @@ match match:
         pass
 match = lambda query: query == event
 print(match(12))
-"#,
-            "<test>",
-        )
-        .unwrap();
-        insta::assert_debug_snapshot!(parse_ast);
+"#;
+    insta::assert_debug_snapshot!(ast::Suite::parse(source, "<test>").unwrap());
     }
 
     #[test]

parser/src/python.lalrpop

@@ -86,6 +86,7 @@ SmallStatement: ast::Stmt = {
     GlobalStatement,
     NonlocalStatement,
     AssertStatement,
+    TypeAliasStatement,
 };
 
 PassStatement: ast::Stmt = {
@@ -979,6 +980,25 @@ FuncDef: ast::Stmt = {
     },
 };
 
+TypeAliasName: ast::Expr = {
+    <location:@L> <name:Identifier> <end_location:@R> => ast::Expr::Name(
+        ast::ExprName { id: name, ctx: ast::ExprContext::Load, range: (location..end_location).into() },
+    ),
+}
+
+TypeAliasStatement: ast::Stmt = {
+    <location:@L> "type" <name:TypeAliasName> <type_params:TypeParamList?> "=" <value:Test<"all">> <end_location:@R> => {
+        ast::Stmt::TypeAlias(
+            ast::StmtTypeAlias {
+                name: Box::new(name),
+                value: Box::new(value),
+                type_params: type_params.unwrap_or_default(),
+                range: (location..end_location).into()
+            },
+        )
+    },
+};
+
 Parameters: ast::Arguments = {
     <location:@L> "(" <a: (ParameterList<TypedParameter, StarTypedParameter, DoubleStarTypedParameter>)?> ")" <end_location:@R> =>? {
         a.as_ref().map(validate_arguments).transpose()?;
@@ -1751,6 +1771,7 @@ extern {
         "raise" => token::Tok::Raise,
         "return" => token::Tok::Return,
         "try" => token::Tok::Try,
+        "type" => token::Tok::Type,
         "while" => token::Tok::While,
         "match" => token::Tok::Match,
         "case" => token::Tok::Case,