Add experimental pyo3-wrapper feature

Author: youknowone

ast/Cargo.toml

@@ -16,6 +16,10 @@ unparse = ["rustpython-literal"]
 visitor = []
 all-nodes-with-ranges = []
 
+# This feature is experimental
+# It reimplements AST types, but currently both slower than python AST types and limited to use in other API
+pyo3-wrapper = ["pyo3"]
+
 [dependencies]
 rustpython-parser-core = { workspace = true }
 rustpython-literal = { workspace = true, optional = true }

ast/asdl_rs.py

@@ -1041,6 +1041,299 @@ def emit_to_pyo3_with_fields(self, cons, name):
                 0,
             )
 
+class Pyo3StructVisitor(EmitVisitor):
+    """Visitor to generate type-defs for AST."""
+
+    def __init__(self, namespace, *args, **kw):
+        self.namespace = namespace
+        self.borrow = True
+        super().__init__(*args, **kw)
+
+    @property
+    def generics(self):
+        if self.namespace == "ranged":
+            return "<TextRange>"
+        elif self.namespace == "located":
+            return "<SourceRange>"
+        else:
+            assert False, self.namespace
+
+    @property
+    def module_name(self):
+        name = f"rustpython_ast.{self.namespace}"
+        return name
+
+    @property
+    def ref_def(self):
+        return "&'static " if self.borrow else ""
+
+    @property
+    def ref(self):
+        return "&" if self.borrow else ""
+
+    def emit_class(self, name, rust_name, simple, base="super::AST"):
+        info = self.type_info[name]
+        if simple:
+            generics = ""
+        else:
+            generics = self.generics
+        if info.is_sum:
+            subclass = ", subclass"
+            body = ""
+            into = f"{rust_name}"
+        else:
+            subclass = ""
+            body = f"(pub {self.ref_def} crate::{rust_name}{generics})"
+            into = f"{rust_name}(node)"
+
+        self.emit(
+            textwrap.dedent(
+                f"""
+                #[pyclass(module="{self.module_name}", name="_{name}", extends={base}, frozen{subclass})]
+                #[derive(Clone, Debug)]
+                pub struct {rust_name} {body};
+
+                impl From<{self.ref_def} crate::{rust_name}{generics}> for {rust_name} {{
+                    fn from({"" if body else "_"}node: {self.ref_def} crate::{rust_name}{generics}) -> Self {{
+                        {into}
+                    }}
+                }}
+                """
+            ),
+            0,
+        )
+        if subclass:
+            self.emit(
+                textwrap.dedent(
+                    f"""
+                    #[pymethods]
+                    impl {rust_name} {{
+                        #[new]
+                        fn new() -> PyClassInitializer<Self> {{
+                            PyClassInitializer::from(AST)
+                                .add_subclass(Self)
+                        }}
+
+                    }}
+                    impl ToPyObject for {rust_name} {{
+                        fn to_object(&self, py: Python) -> PyObject {{
+                            let initializer = PyClassInitializer::from(AST)
+                            .add_subclass(self.clone());
+                            Py::new(py, initializer).unwrap().into_py(py)
+                        }}
+                    }}
+                    """
+                ),
+                0,
+            )
+        else:
+            if base != "super::AST":
+                add_subclass = f".add_subclass({base})"
+            else:
+                add_subclass = ""
+            self.emit(
+                textwrap.dedent(
+                    f"""
+                    impl ToPyObject for {rust_name} {{
+                        fn to_object(&self, py: Python) -> PyObject {{
+                            let initializer = PyClassInitializer::from(AST)
+                            {add_subclass}
+                            .add_subclass(self.clone());
+                            Py::new(py, initializer).unwrap().into_py(py)
+                        }}
+                    }}
+                    """
+                ),
+                0,
+            )
+
+        if not subclass:
+            self.emit_wrapper(rust_name)
+
+    def emit_getter(self, owner, type_name):
+        self.emit(
+            textwrap.dedent(
+                f"""
+                #[pymethods]
+                impl {type_name} {{
+                """
+            ),
+            0,
+        )
+
+        for field in owner.fields:
+            self.emit(
+                textwrap.dedent(
+                    f"""
+                    #[getter]
+                    #[inline]
+                    fn get_{field.name}(&self, py: Python) -> PyResult<PyObject> {{
+                        self.0.{rust_field(field.name)}.to_pyo3_wrapper(py)
+                    }}
+                    """
+                ),
+                3,
+            )
+
+        self.emit(
+            textwrap.dedent(
+                """
+                }
+                """
+            ),
+            0,
+        )
+
+    def emit_getattr(self, owner, type_name):
+        self.emit(
+            textwrap.dedent(
+                f"""
+                #[pymethods]
+                impl {type_name} {{
+                    fn __getattr__(&self, py: Python, key: &str) -> PyResult<PyObject> {{
+                        let object: Py<PyAny> = match key {{
+                """
+            ),
+            0,
+        )
+
+        for field in owner.fields:
+            self.emit(
+                f'"{field.name}" => self.0.{rust_field(field.name)}.to_pyo3_wrapper(py)?,',
+                3,
+            )
+
+        self.emit(
+            textwrap.dedent(
+                """
+                            _ => todo!(),
+                        };
+                        Ok(object)
+                    }
+                }
+                """
+            ),
+            0,
+        )
+
+    def emit_wrapper(self, rust_name):
+        self.emit(
+            f"""
+            impl ToPyo3Wrapper for crate::{rust_name}{self.generics} {{
+                #[inline]
+                fn to_pyo3_wrapper(&'static self, py: Python) -> PyResult<Py<PyAny>> {{
+                    Ok({rust_name}(self).to_object(py))
+                }}
+            }}
+            """,
+            0,
+        )
+
+    def visitModule(self, mod):
+        for dfn in mod.dfns:
+            self.visit(dfn)
+
+    def visitType(self, type, depth=0):
+        self.visit(type.value, type, depth)
+
+    def visitSum(self, sum, type, depth=0):
+        rust_name = rust_type_name(type.name)
+
+        simple = is_simple(sum)
+        self.emit_class(type.name, rust_name, simple)
+
+        if not simple:
+            self.emit(
+                f"""
+                impl ToPyo3Wrapper for crate::{rust_name}{self.generics} {{
+                    #[inline]
+                    fn to_pyo3_wrapper(&'static self, py: Python) -> PyResult<Py<PyAny>> {{
+                        match &self {{
+                """,
+                0,
+            )
+
+            for cons in sum.types:
+                self.emit(f"Self::{cons.name}(cons) => cons.to_pyo3_wrapper(py),", 3)
+
+            self.emit(
+                """
+                        }
+                    }
+                }
+                """,
+                0,
+            )
+
+        for cons in sum.types:
+            self.visit(cons, rust_name, simple, depth + 1)
+
+    def visitProduct(self, product, type, depth=0):
+        rust_name = rust_type_name(type.name)
+        self.emit_class(type.name, rust_name, False)
+        if self.borrow:
+            self.emit_getter(product, rust_name)
+
+    def visitConstructor(self, cons, parent, simple, depth):
+        if simple:
+            self.emit(
+                f"""
+#[pyclass(module="{self.module_name}", name="_{cons.name}", extends={parent})]
+pub struct {parent}{cons.name};
+
+impl ToPyObject for {parent}{cons.name} {{
+    fn to_object(&self, py: Python) -> PyObject {{
+        let initializer = PyClassInitializer::from(AST)
+        .add_subclass({parent})
+        .add_subclass(Self);
+        Py::new(py, initializer).unwrap().into_py(py)
+    }}
+}}
+                """,
+                depth,
+            )
+        else:
+            self.emit_class(
+                cons.name,
+                f"{parent}{cons.name}",
+                simple=False,
+                base=parent,
+            )
+            if self.borrow:
+                self.emit_getter(cons, f"{parent}{cons.name}")
+
+
+class Pyo3PymoduleVisitor(EmitVisitor):
+    def __init__(self, namespace, *args, **kw):
+        self.namespace = namespace
+        super().__init__(*args, **kw)
+
+    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_fields(name, rust_name, False)
+
+    def visitSum(self, sum, name, depth):
+        rust_name = rust_type_name(name)
+        simple = is_simple(sum)
+        self.emit_fields(name, rust_name, True)
+
+        for cons in sum.types:
+            self.visit(cons, name, simple, depth)
+
+    def visitConstructor(self, cons, parent, simple, depth):
+        rust_name = rust_type_name(parent) + rust_type_name(cons.name)
+        self.emit_fields(cons.name, rust_name, simple)
+
+    def emit_fields(self, name, rust_name, simple):
+        self.emit(f"super::init_type::<{rust_name}, crate::generic::{rust_name}>(py, ast_module, m)?;", 1)
+
 
 class StdlibClassDefVisitor(EmitVisitor):
     def visitModule(self, mod):
@@ -1452,6 +1745,60 @@ def write_to_pyo3_simple(type_info, f):
         )
 
 
+def write_pyo3_wrapper(mod, type_info, namespace, f):
+    Pyo3StructVisitor(namespace, f, type_info).visit(mod)
+
+    if namespace == "located":
+        for type_info in type_info.values():
+            if not type_info.is_simple or not type_info.is_sum:
+                continue
+
+            rust_name = type_info.rust_sum_name
+            f.write(
+                f"""
+                impl ToPyo3Wrapper for crate::generic::{rust_name} {{
+                    #[inline]
+                    fn to_pyo3_wrapper(&self, py: Python) -> PyResult<Py<PyAny>> {{
+                        match &self {{
+                """,
+            )
+            for cons in type_info.type.value.types:
+                f.write(
+                    f"Self::{cons.name} => Ok({rust_name}{cons.name}.to_object(py)),",
+                )
+            f.write(
+                """
+                        }
+                    }
+                }
+                """,
+            )
+
+            for cons in type_info.type.value.types:
+                f.write(
+                    f"""
+                    impl ToPyo3Wrapper for crate::generic::{rust_name}{cons.name} {{
+                        #[inline]
+                        fn to_pyo3_wrapper(&self, py: Python) -> PyResult<Py<PyAny>> {{
+                            Ok({rust_name}{cons.name}.to_object(py))
+                        }}
+                    }}
+                    """
+                )
+
+    f.write(
+        """
+        pub fn add_to_module(py: Python, m: &PyModule) -> PyResult<()> {
+            super::init_module(py, m)?;
+
+            let ast_module = PyModule::import(py, "_ast")?;
+        """
+    )
+
+    Pyo3PymoduleVisitor(namespace, f, type_info).visit(mod)
+    f.write("Ok(())\n}")
+
+
 def write_ast_mod(mod, type_info, f):
     f.write(
         textwrap.dedent(
@@ -1498,6 +1845,8 @@ def main(
         ("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)),
+        ("pyo3_wrapper_located", p(write_pyo3_wrapper, mod, type_info, "located")),
+        ("pyo3_wrapper_ranged", p(write_pyo3_wrapper, mod, type_info, "ranged")),
     ]:
         with (ast_dir / f"{filename}.rs").open("w") as f:
             f.write(auto_gen_msg)