to_pyo3_ast to return &'py

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> {",
@@ -1097,7 +1104,7 @@ class ToPyo3AstVisitor(EmitVisitor):
             f"""
             impl ToPyo3Ast for crate::generic::{rust_name}{self.generics} {{
                 #[inline]
-                fn to_pyo3_ast(&self, {"_" if simple else ""}py: Python) -> PyResult<Py<PyAny>> {{
+                fn to_pyo3_ast<'py>(&self, {"_" if simple else ""}py: Python<'py>) -> PyResult<&'py PyAny> {{
                     let instance = match &self {{
             """,
             0,
@@ -1130,7 +1137,7 @@ class ToPyo3AstVisitor(EmitVisitor):
             f"""
             impl ToPyo3Ast for crate::{name}{self.generics} {{
                 #[inline]
-                    fn to_pyo3_ast(&self, py: Python) -> PyResult<Py<PyAny>> {{
+                    fn to_pyo3_ast<'py>(&self, py: Python<'py>) -> PyResult<&'py PyAny> {{
                         let cache = Self::py_type_cache().get().unwrap();
             """,
             0,
@@ -1144,10 +1151,37 @@ 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"{rust_field(field.name)}.to_object(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)?,",
                     3,
@@ -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,
@@ -1858,7 +1892,7 @@ def write_to_pyo3_simple(type_info, f):
             f"""
             impl ToPyo3Ast for crate::generic::{rust_name} {{
                 #[inline]
-                fn to_pyo3_ast(&self, _py: Python) -> PyResult<Py<PyAny>> {{
+                fn to_pyo3_ast<'py>(&self, py: Python<'py>) -> PyResult<&'py PyAny> {{
                     let cell = match &self {{
             """,
         )
@@ -1869,7 +1903,7 @@ def write_to_pyo3_simple(type_info, f):
         f.write(
             """
                     };
-                    Ok(cell.get().unwrap().1.clone())
+                    Ok(Py::<PyAny>::as_ref(&cell.get().unwrap().1, py))
                 }
             }
             """,

ast/src/pyo3.rs

@@ -3,7 +3,8 @@ use num_complex::Complex64;
 use once_cell::sync::OnceCell;
 use pyo3::{
     prelude::*,
-    types::{PyBytes, PyList, PyString, PyTuple},
+    types::{PyBool, PyBytes, PyList, PyString, PyTuple},
+    ToPyObject,
 };
 
 pub trait Pyo3Node {
@@ -16,28 +17,28 @@ pub trait Pyo3Node {
 }
 
 pub trait ToPyo3Ast {
-    fn to_pyo3_ast(&self, py: Python) -> PyResult<Py<PyAny>>;
+    fn to_pyo3_ast<'py>(&self, py: Python<'py>) -> PyResult<&'py PyAny>;
 }
 
 impl<T: ToPyo3Ast> ToPyo3Ast for Box<T> {
     #[inline]
-    fn to_pyo3_ast(&self, py: Python) -> PyResult<Py<PyAny>> {
+    fn to_pyo3_ast<'py>(&self, py: Python<'py>) -> 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>> {
+    fn to_pyo3_ast<'py>(&self, py: Python<'py>) -> PyResult<&'py PyAny> {
         match self {
             Some(ast) => ast.to_pyo3_ast(py),
-            None => Ok(py.None()),
+            None => Ok(ast_cache().none_ref(py)),
         }
     }
 }
 
 impl<T: ToPyo3Ast> ToPyo3Ast for Vec<T> {
-    fn to_pyo3_ast(&self, py: Python) -> PyResult<Py<PyAny>> {
+    fn to_pyo3_ast<'py>(&self, py: Python<'py>) -> PyResult<&'py PyAny> {
         let elts = self
             .iter()
             .map(|item| item.to_pyo3_ast(py))
@@ -49,60 +50,74 @@ impl<T: ToPyo3Ast> ToPyo3Ast for Vec<T> {
 
 impl ToPyo3Ast for crate::Identifier {
     #[inline]
-    fn to_pyo3_ast(&self, py: Python) -> PyResult<Py<PyAny>> {
-        Ok(self.as_str().to_object(py))
+    fn to_pyo3_ast<'py>(&self, py: Python<'py>) -> PyResult<&'py PyAny> {
+        Ok(PyString::new(py, self.as_str()).into())
     }
 }
 
 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))
+    fn to_pyo3_ast<'py>(&self, py: Python<'py>) -> PyResult<&'py PyAny> {
+        Ok(PyString::new(py, self.as_str()).into())
     }
 }
 
 impl ToPyo3Ast for bool {
     #[inline]
-    fn to_pyo3_ast(&self, py: Python) -> PyResult<Py<PyAny>> {
-        Ok((*self as u32).to_object(py))
+    fn to_pyo3_ast<'py>(&self, py: Python<'py>) -> PyResult<&'py PyAny> {
+        Ok(ast_cache().bool_int(py, *self))
     }
 }
 
 impl ToPyo3Ast for ConversionFlag {
     #[inline]
-    fn to_pyo3_ast(&self, py: Python) -> PyResult<Py<PyAny>> {
-        Ok((*self as i8).to_object(py))
+    fn to_pyo3_ast<'py>(&self, py: Python<'py>) -> PyResult<&'py PyAny> {
+        Ok(ast_cache().conversion_flag(py, *self))
     }
 }
 
-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),
+impl ToPyObject for crate::Constant {
+    fn to_object(&self, py: Python) -> PyObject {
+        let cache = ast_cache();
+        match self {
+            crate::Constant::None => cache.none.clone_ref(py),
+            crate::Constant::Bool(bool) => cache.bool(py, *bool).into(),
             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()
+                let elts: Vec<_> = elts.iter().map(|c| c.to_object(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 ToPyo3Ast for crate::Constant {
+//     #[inline]
+//     fn to_pyo3_ast<'py>(&self, py: Python<'py>) -> PyResult<&'py PyAny> {
+//         let cache = ast_cache();
+//         let value = match self {
+//             crate::Constant::None => cache.none_ref(py),
+//             crate::Constant::Bool(bool) => cache.bool_ref(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;
 
@@ -125,26 +140,70 @@ fn cache_py_type<N: Pyo3Node + Node>(ast_module: &PyAny) -> PyResult<()> {
     Ok(())
 }
 
-struct AstKeyCache {
+// 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>),
 }
 
-fn ast_key_cache() -> &'static OnceCell<AstKeyCache> {
+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<AstKeyCache> = OnceCell::new();
+        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_key_cache().get_or_init(|| AstKeyCache {
+    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)