Avoids panic on class initialization when adding a class to the module

PyO3 · Feb 21, 2024 · dc81d76 · dc81d76
1 parent 3796ffb
commit dc81d76
Show file tree

Hide file tree

Showing 5 changed files with 20 additions and 4 deletions.
diff --git a/pyo3-macros-backend/src/pyclass.rs b/pyo3-macros-backend/src/pyclass.rs
@@ -1290,6 +1290,16 @@ fn impl_pytypeinfo(
                     .get_or_init(py)
                     .as_type_ptr()
             }
+
+            #[inline]
+            fn try_type_object_bound(py: _pyo3::Python<'_>) -> _pyo3::PyResult<_pyo3::Bound<'_, _pyo3::types::PyType>> {
+                use ::std::borrow::ToOwned;
+                #deprecations
+
+                <#cls as _pyo3::impl_::pyclass::PyClassImpl>::lazy_type_object()
+                    .get_or_try_init(py)
+                    .map(|t| t.to_owned())
+            }
         }
     }
 }

diff --git a/src/impl_/pyclass/lazy_type_object.rs b/src/impl_/pyclass/lazy_type_object.rs
@@ -54,7 +54,7 @@ impl<T: PyClass> LazyTypeObject<T> {
     }
 
     /// Fallible version of the above.
-    pub(crate) fn get_or_try_init<'py>(&self, py: Python<'py>) -> PyResult<&Bound<'py, PyType>> {
+    pub fn get_or_try_init<'py>(&self, py: Python<'py>) -> PyResult<&Bound<'py, PyType>> {
         self.0
             .get_or_try_init(py, create_type_object::<T>, T::NAME, T::items_iter())
     }

diff --git a/src/impl_/pymodule.rs b/src/impl_/pymodule.rs
@@ -143,7 +143,7 @@ pub trait PyAddToModule {
 
 impl<T: PyTypeInfo> PyAddToModule for T {
     fn add_to_module(module: &Bound<'_, PyModule>) -> PyResult<()> {
-        module.add(Self::NAME, Self::type_object_bound(module.py()))
+        module.add(Self::NAME, Self::try_type_object_bound(module.py())?)
     }
 }
 

diff --git a/src/type_object.rs b/src/type_object.rs
@@ -3,7 +3,7 @@
 use crate::ffi_ptr_ext::FfiPtrExt;
 use crate::types::any::PyAnyMethods;
 use crate::types::{PyAny, PyType};
-use crate::{ffi, Bound, PyNativeType, Python};
+use crate::{ffi, Bound, PyNativeType, PyResult, Python};
 
 /// `T: PyLayout<U>` represents that `T` is a concrete representation of `U` in the Python heap.
 /// E.g., `PyCell` is a concrete representation of all `pyclass`es, and `ffi::PyObject`
@@ -96,6 +96,12 @@ pub unsafe trait PyTypeInfo: Sized + HasPyGilRef {
         }
     }
 
+    /// Returns the safe abstraction over the type object or an error if initialization fails
+    #[inline]
+    fn try_type_object_bound(py: Python<'_>) -> PyResult<Bound<'_, PyType>> {
+        Ok(Self::type_object_bound(py))
+    }
+
     /// Checks if `object` is an instance of this type or a subclass of this type.
     #[inline]
     #[cfg_attr(

diff --git a/src/types/module.rs b/src/types/module.rs
@@ -633,7 +633,7 @@ impl<'py> PyModuleMethods<'py> for Bound<'py, PyModule> {
         T: PyClass,
     {
         let py = self.py();
-        self.add(T::NAME, T::lazy_type_object().get_or_try_init(py)?)
+        self.add(T::NAME, T::try_type_object_bound(py)?)
     }
 
     fn add_wrapped<T>(&self, wrapper: &impl Fn(Python<'py>) -> T) -> PyResult<()>