rustpython_ast + pyo3 (#25)

Cargo.toml

@@ -30,6 +30,7 @@ log = "0.4.16"
 num-complex = "0.4.0"
 num-bigint = "0.4.3"
 num-traits = "0.2"
+pyo3 = { version = "0.18.3" }
 rand = "0.8.5"
 serde = "1.0"
 static_assertions = "1.1"

ast/Cargo.toml

@@ -23,3 +23,6 @@ rustpython-literal = { workspace = true, optional = true }
 is-macro = { workspace = true }
 num-bigint = { workspace = true }
 static_assertions = "1.1.0"
+num-complex = { workspace = true }
+once_cell = { workspace = true }
+pyo3 = { workspace = true, optional = true, features = ["num-bigint", "num-complex"] }

ast/asdl_rs.py

@@ -937,6 +937,111 @@ class LocatedDefVisitor(EmitVisitor):
         )
 
 
+class ToPyo3AstVisitor(EmitVisitor):
+    """Visitor to generate type-defs for AST."""
+
+    def __init__(self, namespace, *args, **kw):
+        super().__init__(*args, **kw)
+        self.namespace = namespace
+
+    @property
+    def generics(self):
+        if self.namespace == "ranged":
+            return "<TextRange>"
+        elif self.namespace == "located":
+            return "<SourceRange>"
+        else:
+            assert False, self.namespace
+
+    def visitModule(self, mod):
+        for dfn in mod.dfns:
+            self.visit(dfn)
+
+    def visitType(self, type, depth=0):
+        self.visit(type.value, type.name, depth)
+
+    def visitProduct(self, product, name, depth=0):
+        rust_name = rust_type_name(name)
+        self.emit_to_pyo3_with_fields(product, rust_name)
+
+    def visitSum(self, sum, name, depth=0):
+        rust_name = rust_type_name(name)
+        simple = is_simple(sum)
+        if is_simple(sum):
+            return
+
+        self.emit(
+            f"""
+            impl ToPyo3Ast for crate::generic::{rust_name}{self.generics} {{
+                #[inline]
+                fn to_pyo3_ast(&self, {"_" if simple else ""}py: Python) -> PyResult<Py<PyAny>> {{
+                    let instance = match &self {{
+                        """,
+            0,
+        )
+        for cons in sum.types:
+            self.emit(
+                f"""crate::{rust_name}::{cons.name}(cons) => cons.to_pyo3_ast(py)?,""",
+                depth,
+            )
+        self.emit(
+            """
+                    };
+                    Ok(instance)
+                }
+            }
+            """,
+            0,
+        )
+
+        for cons in sum.types:
+            self.visit(cons, rust_name, depth)
+
+    def visitConstructor(self, cons, parent, depth):
+        self.emit_to_pyo3_with_fields(cons, f"{parent}{cons.name}")
+
+    def emit_to_pyo3_with_fields(self, cons, name):
+        if cons.fields:
+            self.emit(
+                f"""
+                impl ToPyo3Ast for crate::{name}{self.generics} {{
+                    #[inline]
+                    fn to_pyo3_ast(&self, py: Python) -> PyResult<Py<PyAny>> {{
+                        let cache = Self::py_type_cache().get().unwrap();
+                        let instance = cache.0.call1(py, (
+                """,
+                0,
+            )
+            for field in cons.fields:
+                self.emit(
+                    f"self.{rust_field(field.name)}.to_pyo3_ast(py)?,",
+                    3,
+                )
+            self.emit(
+                """
+                        ))?;
+                        Ok(instance)
+                    }
+                }
+                """,
+                0,
+            )
+        else:
+            self.emit(
+                f"""
+                impl ToPyo3Ast for crate::{name}{self.generics} {{
+                    #[inline]
+                    fn to_pyo3_ast(&self, py: Python) -> PyResult<Py<PyAny>> {{
+                        let cache = Self::py_type_cache().get().unwrap();
+                        let instance = cache.0.call0(py)?;
+                        Ok(instance)
+                    }}
+                }}
+                """,
+                0,
+            )
+
+
 class StdlibClassDefVisitor(EmitVisitor):
     def visitModule(self, mod):
         for dfn in mod.dfns:
@@ -1271,6 +1376,82 @@ def write_located_def(mod, type_info, f):
     LocatedDefVisitor(f, type_info).visit(mod)
 
 
+def write_pyo3_node(type_info, f):
+    def write(info: TypeInfo):
+        rust_name = info.rust_sum_name
+        if info.is_simple:
+            generics = ""
+        else:
+            generics = "<R>"
+
+        f.write(
+            textwrap.dedent(
+                f"""
+                impl{generics} Pyo3Node for crate::generic::{rust_name}{generics} {{
+                    #[inline]
+                    fn py_type_cache() -> &'static OnceCell<(Py<PyAny>, Py<PyAny>)> {{
+                        static PY_TYPE: OnceCell<(Py<PyAny>, Py<PyAny>)> = OnceCell::new();
+                        &PY_TYPE
+                    }}
+                }}
+                """
+            ),
+        )
+
+    for info in type_info.values():
+        write(info)
+
+
+def write_to_pyo3(mod, type_info, f):
+    write_pyo3_node(type_info, f)
+    write_to_pyo3_simple(type_info, f)
+
+    for namespace in ("ranged", "located"):
+        ToPyo3AstVisitor(namespace, f, type_info).visit(mod)
+
+    f.write(
+        """
+        pub fn init(py: Python) -> PyResult<()> {
+            let ast_module = PyModule::import(py, "_ast")?;
+        """
+    )
+
+    for info in type_info.values():
+        rust_name = info.rust_sum_name
+        f.write(f"cache_py_type::<crate::generic::{rust_name}>(ast_module)?;\n")
+    f.write("Ok(())\n}")
+
+
+def write_to_pyo3_simple(type_info, f):
+    for type_info in type_info.values():
+        if not type_info.is_sum:
+            continue
+        if not type_info.is_simple:
+            continue
+
+        rust_name = type_info.rust_sum_name
+        f.write(
+            f"""
+            impl ToPyo3Ast for crate::generic::{rust_name} {{
+                #[inline]
+                fn to_pyo3_ast(&self, _py: Python) -> PyResult<Py<PyAny>> {{
+                    let cell = match &self {{
+            """,
+        )
+        for cons in type_info.type.value.types:
+            f.write(
+                f"""crate::{rust_name}::{cons.name} => crate::{rust_name}{cons.name}::py_type_cache(),""",
+            )
+        f.write(
+            """
+                    };
+                    Ok(cell.get().unwrap().1.clone())
+                }
+            }
+            """,
+        )
+
+
 def write_ast_mod(mod, type_info, f):
     f.write(
         textwrap.dedent(
@@ -1316,6 +1497,7 @@ def main(
         ("ranged", p(write_ranged_def, mod, type_info)),
         ("located", p(write_located_def, mod, type_info)),
         ("visitor", p(write_visitor_def, mod, type_info)),
+        ("to_pyo3", p(write_to_pyo3, mod, type_info)),
     ]:
         with (ast_dir / f"{filename}.rs").open("w") as f:
             f.write(auto_gen_msg)

ast/src/lib.rs

@@ -47,3 +47,6 @@ pub use visitor::Visitor;
 mod optimizer;
 #[cfg(feature = "constant-optimization")]
 pub use optimizer::ConstantOptimizer;
+
+#[cfg(feature = "pyo3")]
+pub mod pyo3;

ast/src/pyo3.rs

@@ -0,0 +1,116 @@
+use crate::{source_code::SourceRange, text_size::TextRange, Node};
+use num_complex::Complex64;
+use once_cell::sync::OnceCell;
+use pyo3::prelude::*;
+use pyo3::types::{PyBytes, PyList, PyTuple};
+
+pub trait Pyo3Node {
+    fn py_type_cache() -> &'static OnceCell<(Py<PyAny>, Py<PyAny>)> {
+        {
+            static PY_TYPE: OnceCell<(Py<PyAny>, Py<PyAny>)> = OnceCell::new();
+            &PY_TYPE
+        }
+    }
+}
+
+pub trait ToPyo3Ast {
+    fn to_pyo3_ast(&self, py: Python) -> PyResult<Py<PyAny>>;
+}
+
+impl<T: ToPyo3Ast> ToPyo3Ast for Box<T> {
+    #[inline]
+    fn to_pyo3_ast(&self, py: Python) -> PyResult<Py<PyAny>> {
+        (**self).to_pyo3_ast(py)
+    }
+}
+
+impl<T: ToPyo3Ast> ToPyo3Ast for Option<T> {
+    #[inline]
+    fn to_pyo3_ast(&self, py: Python) -> PyResult<Py<PyAny>> {
+        match self {
+            Some(ast) => ast.to_pyo3_ast(py),
+            None => Ok(py.None()),
+        }
+    }
+}
+
+impl<T: ToPyo3Ast> ToPyo3Ast for Vec<T> {
+    fn to_pyo3_ast(&self, py: Python) -> PyResult<Py<PyAny>> {
+        let list = PyList::empty(py);
+        for item in self {
+            let py_item = item.to_pyo3_ast(py)?;
+            list.append(py_item)?;
+        }
+        Ok(list.into())
+    }
+}
+
+impl ToPyo3Ast for crate::Identifier {
+    #[inline]
+    fn to_pyo3_ast(&self, py: Python) -> PyResult<Py<PyAny>> {
+        Ok(self.as_str().to_object(py))
+    }
+}
+
+impl ToPyo3Ast for crate::String {
+    #[inline]
+    fn to_pyo3_ast(&self, py: Python) -> PyResult<Py<PyAny>> {
+        Ok(self.as_str().to_object(py))
+    }
+}
+
+impl ToPyo3Ast for crate::Int {
+    #[inline]
+    fn to_pyo3_ast(&self, py: Python) -> PyResult<Py<PyAny>> {
+        Ok((self.to_u32()).to_object(py))
+    }
+}
+
+impl ToPyo3Ast for bool {
+    #[inline]
+    fn to_pyo3_ast(&self, py: Python) -> PyResult<Py<PyAny>> {
+        Ok((*self as u32).to_object(py))
+    }
+}
+
+impl ToPyo3Ast for crate::Constant {
+    #[inline]
+    fn to_pyo3_ast(&self, py: Python) -> PyResult<Py<PyAny>> {
+        let value = match self {
+            crate::Constant::None => py.None(),
+            crate::Constant::Bool(bool) => bool.to_object(py),
+            crate::Constant::Str(string) => string.to_object(py),
+            crate::Constant::Bytes(bytes) => PyBytes::new(py, bytes).into(),
+            crate::Constant::Int(int) => int.to_object(py),
+            crate::Constant::Tuple(elts) => {
+                let elts: PyResult<Vec<_>> = elts.iter().map(|c| c.to_pyo3_ast(py)).collect();
+                PyTuple::new(py, elts?).into()
+            }
+            crate::Constant::Float(f64) => f64.to_object(py),
+            crate::Constant::Complex { real, imag } => Complex64::new(*real, *imag).to_object(py),
+            crate::Constant::Ellipsis => py.Ellipsis(),
+        };
+        Ok(value)
+    }
+}
+
+#[pyclass(module = "rustpython_ast", subclass)]
+pub struct AST;
+
+#[pymethods]
+impl AST {
+    #[new]
+    fn new() -> Self {
+        Self
+    }
+}
+
+fn cache_py_type<N: Pyo3Node + Node>(ast_module: &PyAny) -> PyResult<()> {
+    let class = ast_module.getattr(N::NAME).unwrap();
+    let base = class.getattr("__new__").unwrap();
+    N::py_type_cache().get_or_init(|| (class.into(), base.into()));
+
+    Ok(())
+}
+
+include!("gen/to_pyo3.rs");