Merge pull request #65 from youknowone/refactor-pyo3

to_pyo3_ast to return &'py + Separate rustpython_ast_pyo3

Cargo.toml

@@ -12,6 +12,7 @@ include = ["LICENSE", "Cargo.toml", "src/**/*.rs"]
 resolver = "2"
 members = [
     "ast", "core", "format", "literal", "parser",
+    "ast-pyo3",
     "ruff_text_size", "ruff_source_location",
 ]
 

ast-pyo3/Cargo.toml

@@ -0,0 +1,16 @@
+[package]
+name = "rustpython-ast-pyo3"
+version = "0.1.0"
+edition = "2021"
+
+[features]
+# This feature is experimental
+# It reimplements AST types, but currently both slower than python AST types and limited to use in other API
+wrapper = []
+
+[dependencies]
+rustpython-ast = { workspace = true, features = ["location"] }
+num-complex = { workspace = true }
+once_cell = { workspace = true }
+
+pyo3 = { workspace = true, features = ["num-bigint", "num-complex"] }

ast-pyo3/src/lib.rs

@@ -0,0 +1,5 @@
+mod py_ast;
+#[cfg(feature = "wrapper")]
+pub mod wrapper;
+
+pub use py_ast::{init, PyNode, ToPyAst};

ast-pyo3/src/py_ast.rs

@@ -0,0 +1,190 @@
+use num_complex::Complex64;
+use once_cell::sync::OnceCell;
+use pyo3::{
+    prelude::*,
+    types::{PyBool, PyBytes, PyList, PyString, PyTuple},
+    ToPyObject,
+};
+use rustpython_ast::{
+    self as ast, source_code::SourceRange, text_size::TextRange, ConversionFlag, Node,
+};
+
+pub trait PyNode {
+    fn py_type_cache() -> &'static OnceCell<(Py<PyAny>, Py<PyAny>)> {
+        {
+            static PY_TYPE: OnceCell<(Py<PyAny>, Py<PyAny>)> = OnceCell::new();
+            &PY_TYPE
+        }
+    }
+}
+
+pub trait ToPyAst {
+    fn to_py_ast<'py>(&self, py: Python<'py>) -> PyResult<&'py PyAny>;
+}
+
+impl<T: ToPyAst> ToPyAst for Box<T> {
+    #[inline]
+    fn to_py_ast<'py>(&self, py: Python<'py>) -> PyResult<&'py PyAny> {
+        (**self).to_py_ast(py)
+    }
+}
+
+impl<T: ToPyAst> ToPyAst for Option<T> {
+    #[inline]
+    fn to_py_ast<'py>(&self, py: Python<'py>) -> PyResult<&'py PyAny> {
+        match self {
+            Some(ast) => ast.to_py_ast(py),
+            None => Ok(ast_cache().none_ref(py)),
+        }
+    }
+}
+
+impl<T: ToPyAst> ToPyAst for Vec<T> {
+    fn to_py_ast<'py>(&self, py: Python<'py>) -> PyResult<&'py PyAny> {
+        let elts = self
+            .iter()
+            .map(|item| item.to_py_ast(py))
+            .collect::<Result<Vec<_>, _>>()?;
+        let list = PyList::new(py, elts);
+        Ok(list.into())
+    }
+}
+
+impl ToPyAst for ast::Identifier {
+    #[inline]
+    fn to_py_ast<'py>(&self, py: Python<'py>) -> PyResult<&'py PyAny> {
+        Ok(PyString::new(py, self.as_str()).into())
+    }
+}
+
+impl ToPyAst for ast::String {
+    #[inline]
+    fn to_py_ast<'py>(&self, py: Python<'py>) -> PyResult<&'py PyAny> {
+        Ok(PyString::new(py, self.as_str()).into())
+    }
+}
+
+impl ToPyAst for bool {
+    #[inline]
+    fn to_py_ast<'py>(&self, py: Python<'py>) -> PyResult<&'py PyAny> {
+        Ok(ast_cache().bool_int(py, *self))
+    }
+}
+
+impl ToPyAst for ConversionFlag {
+    #[inline]
+    fn to_py_ast<'py>(&self, py: Python<'py>) -> PyResult<&'py PyAny> {
+        Ok(ast_cache().conversion_flag(py, *self))
+    }
+}
+
+fn constant_to_object(constant: &ast::Constant, py: Python) -> PyObject {
+    let cache = ast_cache();
+    match constant {
+        ast::Constant::None => cache.none.clone_ref(py),
+        ast::Constant::Bool(bool) => cache.bool(py, *bool).into(),
+        ast::Constant::Str(string) => string.to_object(py),
+        ast::Constant::Bytes(bytes) => PyBytes::new(py, bytes).into(),
+        ast::Constant::Int(int) => int.to_object(py),
+        ast::Constant::Tuple(elts) => {
+            let elts: Vec<_> = elts.iter().map(|c| constant_to_object(c, py)).collect();
+            PyTuple::new(py, elts).into()
+        }
+        ast::Constant::Float(f64) => f64.to_object(py),
+        ast::Constant::Complex { real, imag } => Complex64::new(*real, *imag).to_object(py),
+        ast::Constant::Ellipsis => py.Ellipsis(),
+    }
+}
+
+#[pyclass(module = "rustpython_ast", subclass)]
+pub struct Ast;
+
+#[pymethods]
+impl Ast {
+    #[new]
+    fn new() -> Self {
+        Self
+    }
+}
+
+fn cache_py_type<N: PyNode + Node>(ast_module: &PyAny) -> PyResult<()> {
+    let class = ast_module.getattr(N::NAME)?;
+    let base = if std::mem::size_of::<N>() == 0 {
+        class.call0()?
+    } else {
+        class.getattr("__new__")?
+    };
+    N::py_type_cache().get_or_init(|| (class.into(), base.into()));
+    Ok(())
+}
+
+// TODO: This cache must be bound to 'py
+struct AstCache {
+    lineno: Py<PyString>,
+    col_offset: Py<PyString>,
+    end_lineno: Py<PyString>,
+    end_col_offset: Py<PyString>,
+    none: Py<PyAny>,
+    bool_values: (Py<PyBool>, Py<PyBool>),
+    bool_int_values: (Py<PyAny>, Py<PyAny>),
+    conversion_flags: (Py<PyAny>, Py<PyAny>, Py<PyAny>, Py<PyAny>),
+}
+
+impl AstCache {
+    #[inline]
+    fn none_ref<'py>(&'static self, py: Python<'py>) -> &'py PyAny {
+        Py::<PyAny>::as_ref(&self.none, py)
+    }
+    #[inline]
+    fn bool_int<'py>(&'static self, py: Python<'py>, value: bool) -> &'py PyAny {
+        let v = &self.bool_int_values;
+        Py::<PyAny>::as_ref(if value { &v.1 } else { &v.0 }, py)
+    }
+    #[inline]
+    fn bool(&'static self, py: Python, value: bool) -> Py<PyBool> {
+        let v = &self.bool_values;
+        (if value { &v.1 } else { &v.0 }).clone_ref(py)
+    }
+    fn conversion_flag<'py>(&'static self, py: Python<'py>, value: ConversionFlag) -> &'py PyAny {
+        let v = &self.conversion_flags;
+        match value {
+            ConversionFlag::None => v.0.as_ref(py),
+            ConversionFlag::Str => v.1.as_ref(py),
+            ConversionFlag::Ascii => v.2.as_ref(py),
+            ConversionFlag::Repr => v.3.as_ref(py),
+        }
+    }
+}
+
+fn ast_cache_cell() -> &'static OnceCell<AstCache> {
+    {
+        static PY_TYPE: OnceCell<AstCache> = OnceCell::new();
+        &PY_TYPE
+    }
+}
+
+fn ast_cache() -> &'static AstCache {
+    ast_cache_cell().get().unwrap()
+}
+
+pub fn init(py: Python) -> PyResult<()> {
+    ast_cache_cell().get_or_init(|| AstCache {
+        lineno: pyo3::intern!(py, "lineno").into_py(py),
+        col_offset: pyo3::intern!(py, "col_offset").into_py(py),
+        end_lineno: pyo3::intern!(py, "end_lineno").into_py(py),
+        end_col_offset: pyo3::intern!(py, "end_col_offset").into_py(py),
+        none: py.None(),
+        bool_values: (PyBool::new(py, false).into(), PyBool::new(py, true).into()),
+        bool_int_values: ((0).to_object(py), (1).to_object(py)),
+        conversion_flags: (
+            (-1).to_object(py),
+            (b's').to_object(py),
+            (b'a').to_object(py),
+            (b'r').to_object(py),
+        ),
+    });
+
+    init_types(py)
+}
+
+include!("gen/to_py_ast.rs");

ast-pyo3/src/wrapper.rs

@@ -0,0 +1,141 @@
+use crate::PyNode;
+use num_complex::Complex64;
+use pyo3::prelude::*;
+use pyo3::types::{PyBytes, PyList, PyTuple};
+use rustpython_ast::{
+    self as ast, source_code::SourceRange, text_size::TextRange, ConversionFlag, Node,
+};
+
+pub trait ToPyWrapper {
+    fn to_py_wrapper(&'static self, py: Python) -> PyResult<Py<PyAny>>;
+}
+
+impl<T: ToPyWrapper> ToPyWrapper for Box<T> {
+    #[inline]
+    fn to_py_wrapper(&'static self, py: Python) -> PyResult<Py<PyAny>> {
+        (**self).to_py_wrapper(py)
+    }
+}
+
+impl<T: ToPyWrapper> ToPyWrapper for Option<T> {
+    #[inline]
+    fn to_py_wrapper(&'static self, py: Python) -> PyResult<Py<PyAny>> {
+        match self {
+            Some(ast) => ast.to_py_wrapper(py),
+            None => Ok(py.None()),
+        }
+    }
+}
+
+impl ToPyWrapper for ast::Identifier {
+    #[inline]
+    fn to_py_wrapper(&'static self, py: Python) -> PyResult<Py<PyAny>> {
+        Ok(self.as_str().to_object(py))
+    }
+}
+
+impl ToPyWrapper for ast::String {
+    #[inline]
+    fn to_py_wrapper(&'static self, py: Python) -> PyResult<Py<PyAny>> {
+        Ok(self.as_str().to_object(py))
+    }
+}
+
+impl ToPyWrapper for ast::Int {
+    #[inline]
+    fn to_py_wrapper(&'static self, py: Python) -> PyResult<Py<PyAny>> {
+        Ok((self.to_u32()).to_object(py))
+    }
+}
+
+impl ToPyWrapper for bool {
+    #[inline]
+    fn to_py_wrapper(&'static self, py: Python) -> PyResult<Py<PyAny>> {
+        Ok((*self as u32).to_object(py))
+    }
+}
+
+impl ToPyWrapper for ConversionFlag {
+    #[inline]
+    fn to_py_wrapper(&'static self, py: Python) -> PyResult<Py<PyAny>> {
+        Ok((*self as i8).to_object(py))
+    }
+}
+
+impl ToPyWrapper for ast::Constant {
+    fn to_py_wrapper(&'static self, py: Python) -> PyResult<Py<PyAny>> {
+        let value = match self {
+            ast::Constant::None => py.None(),
+            ast::Constant::Bool(bool) => bool.to_object(py),
+            ast::Constant::Str(string) => string.to_object(py),
+            ast::Constant::Bytes(bytes) => PyBytes::new(py, bytes).into(),
+            ast::Constant::Int(int) => int.to_object(py),
+            ast::Constant::Tuple(elts) => {
+                let elts: PyResult<Vec<_>> = elts.iter().map(|c| c.to_py_wrapper(py)).collect();
+                PyTuple::new(py, elts?).into()
+            }
+            ast::Constant::Float(f64) => f64.to_object(py),
+            ast::Constant::Complex { real, imag } => Complex64::new(*real, *imag).to_object(py),
+            ast::Constant::Ellipsis => py.Ellipsis(),
+        };
+        Ok(value)
+    }
+}
+
+impl<T: ToPyWrapper> ToPyWrapper for Vec<T> {
+    fn to_py_wrapper(&'static self, py: Python) -> PyResult<Py<PyAny>> {
+        let list = PyList::empty(py);
+        for item in self {
+            let py_item = item.to_py_wrapper(py)?;
+            list.append(py_item)?;
+        }
+        Ok(list.into())
+    }
+}
+
+#[pyclass(module = "rustpython_ast", subclass)]
+pub struct Ast;
+
+#[pymethods]
+impl Ast {
+    #[new]
+    fn new() -> Self {
+        Self
+    }
+}
+
+pub mod located {
+    pub use super::Ast;
+    use super::*;
+    include!("gen/wrapper_located.rs");
+}
+
+pub mod ranged {
+    pub use super::Ast;
+    use super::*;
+    include!("gen/wrapper_ranged.rs");
+}
+
+fn init_type<P: pyo3::PyClass, N: PyNode + Node>(py: Python, m: &PyModule) -> PyResult<()> {
+    m.add_class::<P>()?;
+    let node = m.getattr(P::NAME)?;
+    if P::NAME != N::NAME {
+        // TODO: no idea how to escape rust keyword on #[pyclass]
+        m.setattr(P::NAME, node)?;
+    }
+    let names: Vec<&'static str> = N::FIELD_NAMES.to_vec();
+    let fields = PyTuple::new(py, names);
+    node.setattr("_fields", fields)?;
+    Ok(())
+}
+
+/// A Python module implemented in Rust.
+fn init_module(py: Python, m: &PyModule) -> PyResult<()> {
+    m.add_class::<Ast>()?;
+
+    let ast = m.getattr("AST")?;
+    let fields = PyTuple::empty(py);
+    ast.setattr("_fields", fields)?;
+
+    Ok(())
+}

ast/Cargo.toml

@@ -15,11 +15,6 @@ fold = []
 unparse = ["rustpython-literal"]
 visitor = []
 all-nodes-with-ranges = []
-pyo3 = ["dep:pyo3", "num-complex", "once_cell"]
-
-# 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 }
@@ -28,6 +23,5 @@ rustpython-literal = { workspace = true, optional = true }
 is-macro = { workspace = true }
 num-bigint = { workspace = true }
 static_assertions = "1.1.0"
-num-complex = { workspace = true, optional = true }
-once_cell = { workspace = true, optional = true }
+
 pyo3 = { workspace = true, optional = true, features = ["num-bigint", "num-complex"] }

ast/asdl_rs.py

@@ -300,10 +300,13 @@ class StructVisitor(EmitVisitor):
 
     def visitModule(self, mod):
         self.emit_attrs(0)
-        self.emit("""
+        self.emit(
+            """
         #[derive(is_macro::Is)]
         pub enum Ast<R=TextRange> {
-        """, 0)
+        """,
+            0,
+        )
         for dfn in mod.dfns:
             rust_name = rust_type_name(dfn.name)
             generics = "" if self.type_info[dfn.name].is_simple else "<R>"
@@ -315,23 +318,29 @@ class StructVisitor(EmitVisitor):
                 # "ast_" prefix to everywhere seems less useful.
                 self.emit('#[is(name = "module")]', 1)
             self.emit(f"{rust_name}({rust_name}{generics}),", 1)
-        self.emit("""
-        }
-        impl<R> Node for Ast<R> {
-            const NAME: &'static str = "AST";
-            const FIELD_NAMES: &'static [&'static str] = &[];
-        }
-        """, 0)
+        self.emit(
+            """
+            }
+            impl<R> Node for Ast<R> {
+                const NAME: &'static str = "AST";
+                const FIELD_NAMES: &'static [&'static str] = &[];
+            }
+            """,
+            0,
+        )
         for dfn in mod.dfns:
             rust_name = rust_type_name(dfn.name)
             generics = "" if self.type_info[dfn.name].is_simple else "<R>"
-            self.emit(f"""
+            self.emit(
+                f"""
             impl<R> From<{rust_name}{generics}> for Ast<R> {{
                 fn from(node: {rust_name}{generics}) -> Self {{
                     Ast::{rust_name}(node)
                 }}
             }}
-            """, 0)
+            """,
+                0,
+            )
 
         for dfn in mod.dfns:
             self.visit(dfn)
@@ -663,9 +672,7 @@ class FoldImplVisitor(EmitVisitor):
 
         cons_type_name = f"{enum_name}{cons.name}"
 
-        self.emit(
-            f"impl<T, U> Foldable<T, U> for {cons_type_name}{apply_t} {{", depth
-        )
+        self.emit(f"impl<T, U> Foldable<T, U> for {cons_type_name}{apply_t} {{", depth)
         self.emit(f"type Mapped = {cons_type_name}{apply_u};", depth + 1)
         self.emit(
             "fn fold<F: Fold<T, TargetU = U> + ?Sized>(self, folder: &mut F) -> Result<Self::Mapped, F::Error> {",
@@ -1095,16 +1102,16 @@ class ToPyo3AstVisitor(EmitVisitor):
 
         self.emit(
             f"""
-            impl ToPyo3Ast for crate::generic::{rust_name}{self.generics} {{
+            impl ToPyAst for ast::{rust_name}{self.generics} {{
                 #[inline]
-                fn to_pyo3_ast(&self, {"_" if simple else ""}py: Python) -> PyResult<Py<PyAny>> {{
+                fn to_py_ast<'py>(&self, {"_" if simple else ""}py: Python<'py>) -> 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)?,",
+                f"ast::{rust_name}::{cons.name}(cons) => cons.to_py_ast(py)?,",
                 1,
             )
         self.emit(
@@ -1128,9 +1135,9 @@ class ToPyo3AstVisitor(EmitVisitor):
         type_info = self.type_info[type.name]
         self.emit(
             f"""
-            impl ToPyo3Ast for crate::{name}{self.generics} {{
+            impl ToPyAst for ast::{name}{self.generics} {{
                 #[inline]
-                    fn to_pyo3_ast(&self, py: Python) -> PyResult<Py<PyAny>> {{
+                    fn to_py_ast<'py>(&self, py: Python<'py>) -> PyResult<&'py PyAny> {{
                         let cache = Self::py_type_cache().get().unwrap();
             """,
             0,
@@ -1144,12 +1151,39 @@ class ToPyo3AstVisitor(EmitVisitor):
                 1,
             )
             self.emit(
-                "let instance = cache.0.call1(py, (",
+                """
+                let instance = Py::<PyAny>::as_ref(&cache.0, py).call1((
+                """,
                 1,
             )
             for field in cons.fields:
+                if field.type == "constant":
+                    self.emit(
+                        f"constant_to_object({rust_field(field.name)}, py),",
+                        3,
+                    )
+                    continue
+                if field.type == "int":
+                    if field.name == "level":
+                        assert field.opt
+                        self.emit(
+                            f"{rust_field(field.name)}.map_or_else(|| py.None(), |level| level.to_u32().to_object(py)),",
+                            3,
+                        )
+                        continue
+                    if field.name in (
+                        "lineno",
+                        "col_offset",
+                        "end_lineno",
+                        "end_col_offset",
+                    ):
+                        self.emit(
+                            f"{rust_field(field.name)}.to_u32().to_object(py),",
+                            3,
+                        )
+                        continue
                 self.emit(
-                    f"{rust_field(field.name)}.to_pyo3_ast(py)?,",
+                    f"{rust_field(field.name)}.to_py_ast(py)?,",
                     3,
                 )
             self.emit(
@@ -1158,7 +1192,7 @@ class ToPyo3AstVisitor(EmitVisitor):
             )
         else:
             self.emit(
-                "let instance = cache.0.call0(py)?;",
+                "let instance = Py::<PyAny>::as_ref(&cache.0, py).call0()?;",
                 1,
             )
             self.emit(
@@ -1168,12 +1202,12 @@ class ToPyo3AstVisitor(EmitVisitor):
         if type.value.attributes and self.namespace == "located":
             self.emit(
                 """
-                let cache = ast_key_cache().get().unwrap();
-                instance.setattr(py, cache.lineno.as_ref(py), _range.start.row.get())?;
-                instance.setattr(py, cache.col_offset.as_ref(py), _range.start.column.get())?;
+                let cache = ast_cache();
+                instance.setattr(cache.lineno.as_ref(py), _range.start.row.get())?;
+                instance.setattr(cache.col_offset.as_ref(py), _range.start.column.get())?;
                 if let Some(end) = _range.end {
-                    instance.setattr(py, cache.end_lineno.as_ref(py), end.row.get())?;
-                    instance.setattr(py, cache.end_col_offset.as_ref(py), end.column.get())?;
+                    instance.setattr(cache.end_lineno.as_ref(py), end.row.get())?;
+                    instance.setattr(cache.end_col_offset.as_ref(py), end.column.get())?;
                 }
                 """,
                 1,
@@ -1218,7 +1252,7 @@ class Pyo3StructVisitor(EmitVisitor):
     def ref(self):
         return "&" if self.borrow else ""
 
-    def emit_class(self, name, rust_name, simple, base="super::AST"):
+    def emit_class(self, name, rust_name, simple, base="super::Ast"):
         info = self.type_info[name]
         if simple:
             generics = ""
@@ -1230,7 +1264,7 @@ class Pyo3StructVisitor(EmitVisitor):
             into = f"{rust_name}"
         else:
             subclass = ""
-            body = f"(pub {self.ref_def} crate::{rust_name}{generics})"
+            body = f"(pub {self.ref_def} ast::{rust_name}{generics})"
             into = f"{rust_name}(node)"
 
         self.emit(
@@ -1239,8 +1273,8 @@ class Pyo3StructVisitor(EmitVisitor):
             #[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 {{
+            impl From<{self.ref_def} ast::{rust_name}{generics}> for {rust_name} {{
+                fn from({"" if body else "_"}node: {self.ref_def} ast::{rust_name}{generics}) -> Self {{
                     {into}
                 }}
             }}
@@ -1254,7 +1288,7 @@ class Pyo3StructVisitor(EmitVisitor):
                 impl {rust_name} {{
                     #[new]
                     fn new() -> PyClassInitializer<Self> {{
-                        PyClassInitializer::from(AST)
+                        PyClassInitializer::from(Ast)
                             .add_subclass(Self)
                     }}
 
@@ -1269,7 +1303,7 @@ class Pyo3StructVisitor(EmitVisitor):
                 0,
             )
         else:
-            if base != "super::AST":
+            if base != "super::Ast":
                 add_subclass = f".add_subclass({base})"
             else:
                 add_subclass = ""
@@ -1277,7 +1311,7 @@ class Pyo3StructVisitor(EmitVisitor):
                 f"""
                 impl ToPyObject for {rust_name} {{
                     fn to_object(&self, py: Python) -> PyObject {{
-                        let initializer = PyClassInitializer::from(AST)
+                        let initializer = PyClassInitializer::from(Ast)
                         {add_subclass}
                         .add_subclass(self.clone());
                         Py::new(py, initializer).unwrap().into_py(py)
@@ -1305,7 +1339,7 @@ class Pyo3StructVisitor(EmitVisitor):
                 #[getter]
                 #[inline]
                 fn get_{field.name}(&self, py: Python) -> PyResult<PyObject> {{
-                    self.0.{rust_field(field.name)}.to_pyo3_wrapper(py)
+                    self.0.{rust_field(field.name)}.to_py_wrapper(py)
                 }}
                 """,
                 3,
@@ -1331,7 +1365,7 @@ class Pyo3StructVisitor(EmitVisitor):
 
         for field in owner.fields:
             self.emit(
-                f'"{field.name}" => self.0.{rust_field(field.name)}.to_pyo3_wrapper(py)?,',
+                f'"{field.name}" => self.0.{rust_field(field.name)}.to_py_wrapper(py)?,',
                 3,
             )
 
@@ -1349,9 +1383,9 @@ class Pyo3StructVisitor(EmitVisitor):
     def emit_wrapper(self, rust_name):
         self.emit(
             f"""
-            impl ToPyo3Wrapper for crate::{rust_name}{self.generics} {{
+            impl ToPyWrapper for ast::{rust_name}{self.generics} {{
                 #[inline]
-                fn to_pyo3_wrapper(&'static self, py: Python) -> PyResult<Py<PyAny>> {{
+                fn to_py_wrapper(&'static self, py: Python) -> PyResult<Py<PyAny>> {{
                     Ok({rust_name}(self).to_object(py))
                 }}
             }}
@@ -1375,16 +1409,16 @@ class Pyo3StructVisitor(EmitVisitor):
         if not simple:
             self.emit(
                 f"""
-                impl ToPyo3Wrapper for crate::{rust_name}{self.generics} {{
+                impl ToPyWrapper for ast::{rust_name}{self.generics} {{
                     #[inline]
-                    fn to_pyo3_wrapper(&'static self, py: Python) -> PyResult<Py<PyAny>> {{
+                    fn to_py_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(f"Self::{cons.name}(cons) => cons.to_py_wrapper(py),", 3)
 
             self.emit(
                 """
@@ -1413,7 +1447,7 @@ class Pyo3StructVisitor(EmitVisitor):
 
                 impl ToPyObject for {parent}{cons.name} {{
                     fn to_object(&self, py: Python) -> PyObject {{
-                        let initializer = PyClassInitializer::from(AST)
+                        let initializer = PyClassInitializer::from(Ast)
                         .add_subclass({parent})
                         .add_subclass(Self);
                         Py::new(py, initializer).unwrap().into_py(py)
@@ -1462,9 +1496,7 @@ class Pyo3PymoduleVisitor(EmitVisitor):
         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, m)?;", 1
-        )
+        self.emit(f"super::init_type::<{rust_name}, ast::{rust_name}>(py, m)?;", 1)
 
 
 class StdlibClassDefVisitor(EmitVisitor):
@@ -1812,7 +1844,7 @@ def write_pyo3_node(type_info, f):
 
         f.write(
             f"""
-            impl{generics} Pyo3Node for crate::generic::{rust_name}{generics} {{
+            impl{generics} PyNode for ast::{rust_name}{generics} {{
                 #[inline]
                 fn py_type_cache() -> &'static OnceCell<(Py<PyAny>, Py<PyAny>)> {{
                     static PY_TYPE: OnceCell<(Py<PyAny>, Py<PyAny>)> = OnceCell::new();
@@ -1842,7 +1874,7 @@ def write_to_pyo3(mod, type_info, f):
 
     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(f"cache_py_type::<ast::{rust_name}>(ast_module)?;\n")
     f.write("Ok(())\n}")
 
 
@@ -1856,20 +1888,20 @@ def write_to_pyo3_simple(type_info, f):
         rust_name = type_info.rust_sum_name
         f.write(
             f"""
-            impl ToPyo3Ast for crate::generic::{rust_name} {{
+            impl ToPyAst for ast::{rust_name} {{
                 #[inline]
-                fn to_pyo3_ast(&self, _py: Python) -> PyResult<Py<PyAny>> {{
+                fn to_py_ast<'py>(&self, py: Python<'py>) -> 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"""ast::{rust_name}::{cons.name} => ast::{rust_name}{cons.name}::py_type_cache(),""",
             )
         f.write(
             """
                     };
-                    Ok(cell.get().unwrap().1.clone())
+                    Ok(Py::<PyAny>::as_ref(&cell.get().unwrap().1, py))
                 }
             }
             """,
@@ -1887,9 +1919,9 @@ def write_pyo3_wrapper(mod, type_info, namespace, f):
             rust_name = type_info.rust_sum_name
             f.write(
                 f"""
-                impl ToPyo3Wrapper for crate::generic::{rust_name} {{
+                impl ToPyWrapper for ast::{rust_name} {{
                     #[inline]
-                    fn to_pyo3_wrapper(&self, py: Python) -> PyResult<Py<PyAny>> {{
+                    fn to_py_wrapper(&self, py: Python) -> PyResult<Py<PyAny>> {{
                         match &self {{
                 """,
             )
@@ -1908,9 +1940,9 @@ def write_pyo3_wrapper(mod, type_info, namespace, f):
             for cons in type_info.type.value.types:
                 f.write(
                     f"""
-                    impl ToPyo3Wrapper for crate::generic::{rust_name}{cons.name} {{
+                    impl ToPyWrapper for ast::{rust_name}{cons.name} {{
                         #[inline]
-                        fn to_pyo3_wrapper(&self, py: Python) -> PyResult<Py<PyAny>> {{
+                        fn to_py_wrapper(&self, py: Python) -> PyResult<Py<PyAny>> {{
                             Ok({rust_name}{cons.name}.to_object(py))
                         }}
                     }}
@@ -1949,6 +1981,7 @@ def write_ast_mod(mod, type_info, f):
 def main(
     input_filename,
     ast_dir,
+    ast_pyo3_dir,
     module_filename,
     dump_module=False,
 ):
@@ -1971,14 +2004,20 @@ 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)),
-        ("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)
             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")),
+        ("wrapper_ranged", p(write_pyo3_wrapper, mod, type_info, "ranged")),
+    ]:
+        with (ast_pyo3_dir / f"{filename}.rs").open("w") as f:
+            f.write(auto_gen_msg)
+            write(f)
+
     with module_filename.open("w") as module_file:
         module_file.write(auto_gen_msg)
         write_ast_mod(mod, type_info, module_file)
@@ -1990,6 +2029,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("-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")
 
@@ -1997,6 +2037,7 @@ if __name__ == "__main__":
     main(
         args.input_file,
         args.ast_dir,
+        args.ast_pyo3_dir,
         args.module_file,
         args.dump_module,
     )

ast/src/lib.rs

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

ast/src/pyo3.rs

@@ -1,153 +0,0 @@
-use crate::{source_code::SourceRange, text_size::TextRange, ConversionFlag, Node};
-use num_complex::Complex64;
-use once_cell::sync::OnceCell;
-use pyo3::{
-    prelude::*,
-    types::{PyBytes, PyList, PyString, 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 elts = self
-            .iter()
-            .map(|item| item.to_pyo3_ast(py))
-            .collect::<Result<Vec<_>, _>>()?;
-        let list = PyList::new(py, elts);
-        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 ConversionFlag {
-    #[inline]
-    fn to_pyo3_ast(&self, py: Python) -> PyResult<Py<PyAny>> {
-        Ok((*self as i8).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)?;
-    let base = if std::mem::size_of::<N>() == 0 {
-        class.call0()?
-    } else {
-        class.getattr("__new__")?
-    };
-    N::py_type_cache().get_or_init(|| (class.into(), base.into()));
-    Ok(())
-}
-
-struct AstKeyCache {
-    lineno: Py<PyString>,
-    col_offset: Py<PyString>,
-    end_lineno: Py<PyString>,
-    end_col_offset: Py<PyString>,
-}
-
-fn ast_key_cache() -> &'static OnceCell<AstKeyCache> {
-    {
-        static PY_TYPE: OnceCell<AstKeyCache> = OnceCell::new();
-        &PY_TYPE
-    }
-}
-
-pub fn init(py: Python) -> PyResult<()> {
-    ast_key_cache().get_or_init(|| AstKeyCache {
-        lineno: pyo3::intern!(py, "lineno").into_py(py),
-        col_offset: pyo3::intern!(py, "col_offset").into_py(py),
-        end_lineno: pyo3::intern!(py, "end_lineno").into_py(py),
-        end_col_offset: pyo3::intern!(py, "end_col_offset").into_py(py),
-    });
-
-    init_types(py)
-}
-
-include!("gen/to_pyo3.rs");

ast/src/pyo3_wrapper.rs

@@ -1,128 +0,0 @@
-use crate::pyo3::{Pyo3Node, AST};
-use crate::{source_code::SourceRange, text_size::TextRange, ConversionFlag, Node};
-use num_complex::Complex64;
-use pyo3::prelude::*;
-use pyo3::types::{PyBytes, PyList, PyTuple};
-
-pub trait ToPyo3Wrapper {
-    fn to_pyo3_wrapper(&'static self, py: Python) -> PyResult<Py<PyAny>>;
-}
-
-impl<T: ToPyo3Wrapper> ToPyo3Wrapper for Box<T> {
-    #[inline]
-    fn to_pyo3_wrapper(&'static self, py: Python) -> PyResult<Py<PyAny>> {
-        (**self).to_pyo3_wrapper(py)
-    }
-}
-
-impl<T: ToPyo3Wrapper> ToPyo3Wrapper for Option<T> {
-    #[inline]
-    fn to_pyo3_wrapper(&'static self, py: Python) -> PyResult<Py<PyAny>> {
-        match self {
-            Some(ast) => ast.to_pyo3_wrapper(py),
-            None => Ok(py.None()),
-        }
-    }
-}
-
-impl ToPyo3Wrapper for crate::Identifier {
-    #[inline]
-    fn to_pyo3_wrapper(&'static self, py: Python) -> PyResult<Py<PyAny>> {
-        Ok(self.as_str().to_object(py))
-    }
-}
-
-impl ToPyo3Wrapper for crate::String {
-    #[inline]
-    fn to_pyo3_wrapper(&'static self, py: Python) -> PyResult<Py<PyAny>> {
-        Ok(self.as_str().to_object(py))
-    }
-}
-
-impl ToPyo3Wrapper for crate::Int {
-    #[inline]
-    fn to_pyo3_wrapper(&'static self, py: Python) -> PyResult<Py<PyAny>> {
-        Ok((self.to_u32()).to_object(py))
-    }
-}
-
-impl ToPyo3Wrapper for bool {
-    #[inline]
-    fn to_pyo3_wrapper(&'static self, py: Python) -> PyResult<Py<PyAny>> {
-        Ok((*self as u32).to_object(py))
-    }
-}
-
-impl ToPyo3Wrapper for ConversionFlag {
-    #[inline]
-    fn to_pyo3_wrapper(&'static self, py: Python) -> PyResult<Py<PyAny>> {
-        Ok((*self as i8).to_object(py))
-    }
-}
-
-impl ToPyo3Wrapper for crate::Constant {
-    fn to_pyo3_wrapper(&'static 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_wrapper(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)
-    }
-}
-
-impl<T: ToPyo3Wrapper> ToPyo3Wrapper for Vec<T> {
-    fn to_pyo3_wrapper(&'static self, py: Python) -> PyResult<Py<PyAny>> {
-        let list = PyList::empty(py);
-        for item in self {
-            let py_item = item.to_pyo3_wrapper(py)?;
-            list.append(py_item)?;
-        }
-        Ok(list.into())
-    }
-}
-
-pub mod located {
-    use super::*;
-    pub use crate::pyo3::AST;
-    include!("gen/pyo3_wrapper_located.rs");
-}
-
-pub mod ranged {
-    use super::*;
-    pub use crate::pyo3::AST;
-    include!("gen/pyo3_wrapper_ranged.rs");
-}
-
-fn init_type<P: pyo3::PyClass, N: Pyo3Node + Node>(py: Python, m: &PyModule) -> PyResult<()> {
-    m.add_class::<P>()?;
-    let node = m.getattr(P::NAME)?;
-    if P::NAME != N::NAME {
-        // TODO: no idea how to escape rust keyword on #[pyclass]
-        m.setattr(P::NAME, node)?;
-    }
-    let names: Vec<&'static str> = N::FIELD_NAMES.to_vec();
-    let fields = PyTuple::new(py, names);
-    node.setattr("_fields", fields)?;
-    Ok(())
-}
-
-/// A Python module implemented in Rust.
-fn init_module(py: Python, m: &PyModule) -> PyResult<()> {
-    m.add_class::<AST>()?;
-
-    let ast = m.getattr("AST")?;
-    let fields = PyTuple::empty(py);
-    ast.setattr("_fields", fields)?;
-
-    Ok(())
-}

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/ --module-file ../RustPython/vm/src/stdlib/ast/gen.rs ast/Python.asdl
-rustfmt ast/src/gen/*.rs ../RustPython/vm/src/stdlib/ast/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