Separate rustpython_ast_pyo3

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(())
+}