introduce fast path to LazyTypeObject::get_or_init (#5324)

* introduce fast path to `LazyTypeObject::get_or_init`

* fix changelog name

* reduce duplication

* fixup benchmark

* fixup clippy

* fixup doc

Author: davidhewitt

guide/src/class.md

@@ -1385,7 +1385,12 @@ unsafe impl pyo3::type_object::PyTypeInfo for MyClass {
     #[inline]
     fn type_object_raw(py: pyo3::Python<'_>) -> *mut pyo3::ffi::PyTypeObject {
         <Self as pyo3::impl_::pyclass::PyClassImpl>::lazy_type_object()
-            .get_or_init(py)
+            .get_or_try_init(py)
+            .unwrap_or_else(|e| pyo3::impl_::pyclass::type_object_init_failed(
+                py,
+                e,
+                <Self as pyo3::type_object::PyTypeInfo>::NAME
+            ))
             .as_type_ptr()
     }
 }

newsfragments/5324.changed.md

@@ -0,0 +1 @@
+Add fast-path to `PyTypeInfo::type_object` for `#[pyclass]` types.

pyo3-benches/benches/bench_pyclass.rs

@@ -32,7 +32,7 @@ pub fn first_time_init(b: &mut Bencher<'_>) {
             // This is using an undocumented internal PyO3 API to measure pyclass performance; please
             // don't use this in your own code!
             let ty = LazyTypeObject::<MyClass>::new();
-            ty.get_or_init(py);
+            ty.get_or_try_init(py).unwrap();
         });
     });
 }

pyo3-macros-backend/src/pyclass.rs

@@ -1845,7 +1845,12 @@ fn impl_pytypeinfo(cls: &syn::Ident, attr: &PyClassArgs, ctx: &Ctx) -> TokenStre
             fn type_object_raw(py: #pyo3_path::Python<'_>) -> *mut #pyo3_path::ffi::PyTypeObject {
                 use #pyo3_path::prelude::PyTypeMethods;
                 <#cls as #pyo3_path::impl_::pyclass::PyClassImpl>::lazy_type_object()
-                    .get_or_init(py)
+                    .get_or_try_init(py)
+                    .unwrap_or_else(|e| #pyo3_path::impl_::pyclass::type_object_init_failed(
+                        py,
+                        e,
+                        <Self as #pyo3_path::type_object::PyTypeInfo>::NAME
+                    ))
                     .as_type_ptr()
             }
         }

src/impl_/pyclass.rs

@@ -29,7 +29,7 @@ mod lazy_type_object;
 mod probes;
 
 pub use assertions::*;
-pub use lazy_type_object::LazyTypeObject;
+pub use lazy_type_object::{type_object_init_failed, LazyTypeObject};
 pub use probes::*;
 
 /// Gets the offset of the dictionary from the start of the object in bytes.

src/impl_/pyclass/lazy_type_object.rs

@@ -32,7 +32,7 @@ struct LazyTypeObjectInner {
     // Threads which have begun initialization of the `tp_dict`. Used for
     // reentrant initialization detection.
     initializing_threads: Mutex<Vec<ThreadId>>,
-    tp_dict_filled: GILOnceCell<()>,
+    fully_initialized_type: GILOnceCell<Py<PyType>>,
 }
 
 impl<T> LazyTypeObject<T> {
@@ -43,7 +43,7 @@ impl<T> LazyTypeObject<T> {
             LazyTypeObjectInner {
                 value: GILOnceCell::new(),
                 initializing_threads: Mutex::new(Vec::new()),
-                tp_dict_filled: GILOnceCell::new(),
+                fully_initialized_type: GILOnceCell::new(),
             },
             PhantomData,
         )
@@ -52,15 +52,18 @@ impl<T> LazyTypeObject<T> {
 
 impl<T: PyClass> LazyTypeObject<T> {
     /// Gets the type object contained by this `LazyTypeObject`, initializing it if needed.
-    pub fn get_or_init<'py>(&self, py: Python<'py>) -> &Bound<'py, PyType> {
-        self.get_or_try_init(py).unwrap_or_else(|err| {
-            err.print(py);
-            panic!("failed to create type object for {}", T::NAME)
-        })
+    #[inline]
+    pub fn get_or_try_init<'py>(&self, py: Python<'py>) -> PyResult<&Bound<'py, PyType>> {
+        if let Some(type_object) = self.0.fully_initialized_type.get(py) {
+            // Fast path
+            return Ok(type_object.bind(py));
+        }
+
+        self.try_init(py)
     }
 
-    /// Fallible version of the above.
-    pub(crate) fn get_or_try_init<'py>(&self, py: Python<'py>) -> PyResult<&Bound<'py, PyType>> {
+    #[cold]
+    fn 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())
     }
@@ -116,7 +119,7 @@ impl LazyTypeObjectInner {
         // `tp_dict`, it can still request the type object through `get_or_init`,
         // but the `tp_dict` may appear empty of course.
 
-        if self.tp_dict_filled.get(py).is_some() {
+        if self.fully_initialized_type.get(py).is_some() {
             // `tp_dict` is already filled: ok.
             return Ok(());
         }
@@ -184,8 +187,8 @@ impl LazyTypeObjectInner {
 
         // Now we hold the GIL and we can assume it won't be released until we
         // return from the function.
-        let result = self.tp_dict_filled.get_or_try_init(py, move || {
-            let result = initialize_tp_dict(py, type_object.as_ptr(), items);
+        let result = self.fully_initialized_type.get_or_try_init(py, move || {
+            initialize_tp_dict(py, type_object.as_ptr(), items)?;
             #[cfg(Py_3_14)]
             if is_immutable_type {
                 // freeze immutable types after __dict__ is initialized
@@ -216,13 +219,13 @@ impl LazyTypeObjectInner {
                 self.initializing_threads.lock().unwrap()
             };
             threads.clear();
-            result
+            Ok(type_object.clone().unbind())
         });
 
         if let Err(err) = result {
             return Err(wrap_in_runtime_error(
                 py,
-                err.clone_ref(py),
+                err,
                 format!("An error occurred while initializing `{name}.__dict__`"),
             ));
         }
@@ -249,6 +252,13 @@ fn initialize_tp_dict(
 // This is necessary for making static `LazyTypeObject`s
 unsafe impl<T> Sync for LazyTypeObject<T> {}
 
+/// Used in the macro-expanded implementation of `type_object_raw` for `#[pyclass]` types
+#[cold]
+pub fn type_object_init_failed(py: Python<'_>, err: PyErr, type_name: &str) -> ! {
+    err.write_unraisable(py, None);
+    panic!("failed to create type object for `{type_name}`")
+}
+
 #[cold]
 fn wrap_in_runtime_error(py: Python<'_>, err: PyErr, message: String) -> PyErr {
     let runtime_err = PyRuntimeError::new_err(message);

tests/test_class_attributes.rs

@@ -152,40 +152,10 @@ fn recursive_class_attributes() {
 }
 
 #[test]
+#[cfg(all(Py_3_8, not(Py_GIL_DISABLED)))] // sys.unraisablehook not available until Python 3.8
 fn test_fallible_class_attribute() {
-    use pyo3::{exceptions::PyValueError, types::PyString};
-
-    struct CaptureStdErr<'py> {
-        oldstderr: Bound<'py, PyAny>,
-        string_io: Bound<'py, PyAny>,
-    }
-
-    impl<'py> CaptureStdErr<'py> {
-        fn new(py: Python<'py>) -> PyResult<Self> {
-            let sys = py.import("sys")?;
-            let oldstderr = sys.getattr("stderr")?;
-            let string_io = py.import("io")?.getattr("StringIO")?.call0()?;
-            sys.setattr("stderr", &string_io)?;
-            Ok(Self {
-                oldstderr,
-                string_io,
-            })
-        }
-
-        fn reset(self) -> PyResult<String> {
-            let py = self.string_io.py();
-            let payload = self
-                .string_io
-                .getattr("getvalue")?
-                .call0()?
-                .cast::<PyString>()?
-                .to_cow()?
-                .into_owned();
-            let sys = py.import("sys")?;
-            sys.setattr("stderr", self.oldstderr)?;
-            Ok(payload)
-        }
-    }
+    use common::UnraisableCapture;
+    use pyo3::exceptions::PyValueError;
 
     #[pyclass]
     struct BrokenClass;
@@ -199,21 +169,29 @@ fn test_fallible_class_attribute() {
     }
 
     Python::attach(|py| {
-        let stderr = CaptureStdErr::new(py).unwrap();
+        let capture = UnraisableCapture::install(py);
         assert!(std::panic::catch_unwind(|| py.get_type::<BrokenClass>()).is_err());
-        assert_eq!(
-            stderr.reset().unwrap().trim(),
-            "\
-ValueError: failed to create class attribute
-
-The above exception was the direct cause of the following exception:
 
-RuntimeError: An error occurred while initializing `BrokenClass.fails_to_init`
+        let (err, object) = capture.borrow_mut(py).capture.take().unwrap();
+        assert!(object.is_none(py));
 
-The above exception was the direct cause of the following exception:
+        assert_eq!(
+            err.to_string(),
+            "RuntimeError: An error occurred while initializing class BrokenClass"
+        );
+        let cause = err.cause(py).unwrap();
+        assert_eq!(
+            cause.to_string(),
+            "RuntimeError: An error occurred while initializing `BrokenClass.fails_to_init`"
+        );
+        let cause = cause.cause(py).unwrap();
+        assert_eq!(
+            cause.to_string(),
+            "ValueError: failed to create class attribute"
+        );
+        assert!(cause.cause(py).is_none());
 
-RuntimeError: An error occurred while initializing class BrokenClass"
-        )
+        capture.borrow_mut(py).uninstall(py);
     });
 }