Parse for expr and stmt variants + identifier, constant (#78)

ast/asdl_rs.py

@@ -1951,6 +1951,40 @@ def write_pyo3_wrapper(mod, type_info, namespace, f):
     f.write("Ok(())\n}")
 
 
+def write_parse_def(mod, type_info, f):
+    for info in type_info.values():
+        if info.enum_name not in ["expr", "stmt"]:
+            continue
+
+        type_name = rust_type_name(info.enum_name)
+        cons_name = rust_type_name(info.name)
+
+        f.write(f"""
+        impl Parse for ast::{info.rust_sum_name} {{
+            fn lex_starts_at(
+                source: &str,
+                offset: TextSize,
+            ) -> SoftKeywordTransformer<Lexer<std::str::Chars>> {{
+                ast::{type_name}::lex_starts_at(source, offset)
+            }}
+            fn parse_tokens(
+                lxr: impl IntoIterator<Item = LexResult>,
+                source_path: &str,
+            ) -> Result<Self, ParseError> {{
+                let node = ast::{type_name}::parse_tokens(lxr, source_path)?;
+                match node {{
+                    ast::{type_name}::{cons_name}(node) => Ok(node),
+                    node => Err(ParseError {{
+                        error: ParseErrorType::InvalidToken,
+                        offset: node.range().start(),
+                        source_path: source_path.to_owned(),
+                    }}),
+                }}
+            }}
+        }}
+        """)
+
+
 def write_ast_mod(mod, type_info, f):
     f.write(
         """
@@ -1972,6 +2006,7 @@ def write_ast_mod(mod, type_info, f):
 def main(
     input_filename,
     ast_dir,
+    parser_dir,
     ast_pyo3_dir,
     module_filename,
     dump_module=False,
@@ -2000,6 +2035,14 @@ def main(
             f.write(auto_gen_msg)
             write(f)
 
+    for filename, write in [
+        ("parse", p(write_parse_def, mod, type_info)),
+    ]:
+        with (parser_dir / f"{filename}.rs").open("w") as f:
+            f.write(auto_gen_msg)
+            write(f)
+
+
     for filename, write in [
         ("to_py_ast", p(write_to_pyo3, mod, type_info)),
         ("wrapper_located", p(write_pyo3_wrapper, mod, type_info, "located")),
@@ -2020,6 +2063,7 @@ if __name__ == "__main__":
     parser = ArgumentParser()
     parser.add_argument("input_file", type=Path)
     parser.add_argument("-A", "--ast-dir", type=Path, required=True)
+    parser.add_argument("-P", "--parser-dir", type=Path, required=True)
     parser.add_argument("-O", "--ast-pyo3-dir", type=Path, required=True)
     parser.add_argument("-M", "--module-file", type=Path, required=True)
     parser.add_argument("-d", "--dump-module", action="store_true")
@@ -2028,6 +2072,7 @@ if __name__ == "__main__":
     main(
         args.input_file,
         args.ast_dir,
+        args.parser_dir,
         args.ast_pyo3_dir,
         args.module_file,
         args.dump_module,

parser/src/parser.rs

@@ -59,6 +59,9 @@ where
     fn parse(source: &str, source_path: &str) -> Result<Self, ParseError> {
         Self::parse_starts_at(source, source_path, TextSize::default())
     }
+    fn parse_without_path(source: &str) -> Result<Self, ParseError> {
+        Self::parse(source, "<unknown>")
+    }
     fn parse_starts_at(
         source: &str,
         source_path: &str,
@@ -197,6 +200,52 @@ impl Parse for ast::Expr {
     }
 }
 
+impl Parse for ast::Identifier {
+    fn lex_starts_at(
+        source: &str,
+        offset: TextSize,
+    ) -> SoftKeywordTransformer<Lexer<std::str::Chars>> {
+        ast::Expr::lex_starts_at(source, offset)
+    }
+    fn parse_tokens(
+        lxr: impl IntoIterator<Item = LexResult>,
+        source_path: &str,
+    ) -> Result<Self, ParseError> {
+        let expr = ast::Expr::parse_tokens(lxr, source_path)?;
+        match expr {
+            ast::Expr::Name(name) => Ok(name.id),
+            expr => Err(ParseError {
+                error: ParseErrorType::InvalidToken,
+                offset: expr.range().start(),
+                source_path: source_path.to_owned(),
+            }),
+        }
+    }
+}
+
+impl Parse for ast::Constant {
+    fn lex_starts_at(
+        source: &str,
+        offset: TextSize,
+    ) -> SoftKeywordTransformer<Lexer<std::str::Chars>> {
+        ast::Expr::lex_starts_at(source, offset)
+    }
+    fn parse_tokens(
+        lxr: impl IntoIterator<Item = LexResult>,
+        source_path: &str,
+    ) -> Result<Self, ParseError> {
+        let expr = ast::Expr::parse_tokens(lxr, source_path)?;
+        match expr {
+            ast::Expr::Constant(c) => Ok(c.value),
+            expr => Err(ParseError {
+                error: ParseErrorType::InvalidToken,
+                offset: expr.range().start(),
+                source_path: source_path.to_owned(),
+            }),
+        }
+    }
+}
+
 /// Parse a full Python program usually consisting of multiple lines.
 ///  
 /// This is a convenience function that can be used to parse a full Python program without having to
@@ -505,6 +554,8 @@ pub(super) fn optional_range(start: TextSize, end: TextSize) -> OptionalRange<Te
     OptionalRange::<TextRange>::new(start, end)
 }
 
+include!("gen/parse.rs");
+
 #[cfg(test)]
 mod tests {
     use super::*;
@@ -1033,4 +1084,21 @@ def args_to_tuple(*args: *Ts) -> Tuple[*Ts]: ...
         .unwrap();
         insta::assert_debug_snapshot!(parse_ast);
     }
+
+    #[test]
+    fn test_parse_constant() {
+        use num_traits::ToPrimitive;
+
+        let c = ast::Constant::parse_without_path("'string'").unwrap();
+        assert_eq!(c.str().unwrap(), "string");
+
+        let c = ast::Constant::parse_without_path("10").unwrap();
+        assert_eq!(c.int().unwrap().to_i32().unwrap(), 10);
+    }
+
+    #[test]
+    fn test_parse_identifier() {
+        let i = ast::Identifier::parse_without_path("test").unwrap();
+        assert_eq!(i.as_str(), "test");
+    }
 }

scripts/update_asdl.sh

@@ -4,5 +4,5 @@ set -e
 cd "$(dirname "$(dirname "$0")")"
 
 # rm ast/src/gen/*.rs
-python ast/asdl_rs.py --ast-dir ast/src/gen/ --ast-pyo3-dir ast-pyo3/src/gen/ --module-file ../RustPython/vm/src/stdlib/ast/gen.rs ast/Python.asdl
-rustfmt ast/src/gen/*.rs ast-pyo3/src/gen/*.rs ../RustPython/vm/src/stdlib/ast/gen.rs
+python ast/asdl_rs.py --ast-dir ast/src/gen/ --parser-dir parser/src/gen/ --ast-pyo3-dir ast-pyo3/src/gen/ --module-file ../RustPython/vm/src/stdlib/ast/gen.rs ast/Python.asdl
+rustfmt ast/src/gen/*.rs parser/src/gen/*.rs ast-pyo3/src/gen/*.rs ../RustPython/vm/src/stdlib/ast/gen.rs