Add experimental pyo3-wrapper feature (#41)

* Fix pyo3 unit type value error

* Add experimental pyo3-wrapper feature

* location support

ast/src/lib.rs

@@ -50,3 +50,5 @@ pub use optimizer::ConstantOptimizer;
 
 #[cfg(feature = "pyo3")]
 pub mod pyo3;
+#[cfg(feature = "pyo3-wrapper")]
+pub mod pyo3_wrapper;

ast/src/pyo3.rs

@@ -1,8 +1,10 @@
 use crate::{source_code::SourceRange, text_size::TextRange, ConversionFlag, Node};
 use num_complex::Complex64;
 use once_cell::sync::OnceCell;
-use pyo3::prelude::*;
-use pyo3::types::{PyBytes, PyList, PyTuple};
+use pyo3::{
+    prelude::*,
+    types::{PyBytes, PyList, PyString, PyTuple},
+};
 
 pub trait Pyo3Node {
     fn py_type_cache() -> &'static OnceCell<(Py<PyAny>, Py<PyAny>)> {
@@ -113,11 +115,39 @@ impl AST {
 }
 
 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();
+    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

@@ -0,0 +1,128 @@
+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(())
+}