pyo3/pyclass/
guard.rs

1use crate::impl_::pycell::{PyClassObject, PyClassObjectLayout as _};
2#[cfg(feature = "experimental-inspect")]
3use crate::inspect::TypeHint;
4use crate::pycell::PyBorrowMutError;
5use crate::pycell::{impl_::PyClassBorrowChecker, PyBorrowError};
6use crate::pyclass::boolean_struct::False;
7use crate::{ffi, Borrowed, CastError, FromPyObject, IntoPyObject, Py, PyClass, PyErr};
8use std::convert::Infallible;
9use std::fmt;
10use std::marker::PhantomData;
11use std::ops::{Deref, DerefMut};
12use std::ptr::NonNull;
13
14/// A wrapper type for an immutably borrowed value from a `PyClass`.
15///
16/// Rust has strict aliasing rules - you can either have any number of immutable
17/// (shared) references or one mutable reference. Python's ownership model is
18/// the complete opposite of that - any Python object can be referenced any
19/// number of times, and mutation is allowed from any reference.
20///
21/// PyO3 deals with these differences by employing the [Interior Mutability]
22/// pattern. This requires that PyO3 enforces the borrowing rules and it has two
23/// mechanisms for doing so:
24/// - Statically it can enforce thread-safe access with the
25///   [`Python<'py>`](crate::Python) token. All Rust code holding that token, or
26///   anything derived from it, can assume that they have safe access to the
27///   Python interpreter's state. For this reason all the native Python objects
28///   can be mutated through shared references.
29/// - However, methods and functions in Rust usually *do* need `&mut`
30///   references. While PyO3 can use the [`Python<'py>`](crate::Python) token to
31///   guarantee thread-safe access to them, it cannot statically guarantee
32///   uniqueness of `&mut` references. As such those references have to be
33///   tracked dynamically at runtime, using [`PyClassGuard`] and
34///   [`PyClassGuardMut`] defined in this module. This works similar to std's
35///   [`RefCell`](std::cell::RefCell) type. Especially when building for
36///   free-threaded Python it gets harder to track which thread borrows which
37///   object at any time. This can lead to method calls failing with
38///   [`PyBorrowError`]. In these cases consider using `frozen` classes together
39///   with Rust interior mutability primitives like [`Mutex`](std::sync::Mutex)
40///   instead of using [`PyClassGuardMut`] to get mutable access.
41///
42/// # Examples
43///
44/// You can use [`PyClassGuard`] as an alternative to a `&self` receiver when
45/// - you need to access the pointer of the `PyClass`, or
46/// - you want to get a super class.
47/// ```
48/// # use pyo3::prelude::*;
49/// # use pyo3::PyClassGuard;
50/// #[pyclass(subclass)]
51/// struct Parent {
52///     basename: &'static str,
53/// }
54///
55/// #[pyclass(extends=Parent)]
56/// struct Child {
57///     name: &'static str,
58///  }
59///
60/// #[pymethods]
61/// impl Child {
62///     #[new]
63///     fn new() -> (Self, Parent) {
64///         (Child { name: "Caterpillar" }, Parent { basename: "Butterfly" })
65///     }
66///
67///     fn format(slf: PyClassGuard<'_, Self>) -> String {
68///         // We can get &Self::BaseType by as_super
69///         let basename = slf.as_super().basename;
70///         format!("{}(base: {})", slf.name, basename)
71///     }
72/// }
73/// # Python::attach(|py| {
74/// #     let sub = Py::new(py, Child::new()).unwrap();
75/// #     pyo3::py_run!(py, sub, "assert sub.format() == 'Caterpillar(base: Butterfly)', sub.format()");
76/// # });
77/// ```
78///
79/// See also [`PyClassGuardMut`] and the [guide] for more information.
80///
81/// [Interior Mutability]:
82///     https://doc.rust-lang.org/book/ch15-05-interior-mutability.html
83///     "RefCell<T> and the Interior Mutability Pattern - The Rust Programming
84///     Language"
85/// [guide]: https://pyo3.rs/latest/class.html#bound-and-interior-mutability
86///     "Bound and interior mutability"
87#[repr(transparent)]
88pub struct PyClassGuard<'a, T: PyClass> {
89    ptr: NonNull<ffi::PyObject>,
90    marker: PhantomData<&'a Py<T>>,
91}
92
93impl<'a, T: PyClass> PyClassGuard<'a, T> {
94    pub(crate) fn try_borrow(obj: &'a Py<T>) -> Result<Self, PyBorrowError> {
95        Self::try_from_class_object(obj.get_class_object())
96    }
97
98    fn try_from_class_object(obj: &'a PyClassObject<T>) -> Result<Self, PyBorrowError> {
99        obj.ensure_threadsafe();
100        obj.borrow_checker().try_borrow().map(|_| Self {
101            ptr: NonNull::from(obj).cast(),
102            marker: PhantomData,
103        })
104    }
105
106    pub(crate) fn as_class_object(&self) -> &'a PyClassObject<T> {
107        // SAFETY: `ptr` by construction points to a `PyClassObject<T>` and is
108        // valid for at least 'a
109        unsafe { self.ptr.cast().as_ref() }
110    }
111
112    /// Consumes the [`PyClassGuard`] and returns a [`PyClassGuardMap`] for a component of the
113    /// borrowed data
114    ///
115    /// # Examples
116    ///
117    /// ```
118    /// # use pyo3::prelude::*;
119    /// # use pyo3::PyClassGuard;
120    ///
121    /// #[pyclass]
122    /// pub struct MyClass {
123    ///     msg: String,
124    /// }
125    ///
126    /// # Python::attach(|py| {
127    /// let obj = Bound::new(py, MyClass { msg: String::from("hello") })?;
128    /// let msg = obj.extract::<PyClassGuard<'_, MyClass>>()?.map(|c| &c.msg);
129    /// assert_eq!(&*msg, "hello");
130    /// # Ok::<_, PyErr>(())
131    /// # }).unwrap();
132    /// ```
133    pub fn map<F, U: ?Sized>(self, f: F) -> PyClassGuardMap<'a, U, false>
134    where
135        F: FnOnce(&T) -> &U,
136    {
137        let slf = std::mem::ManuallyDrop::new(self); // the borrow is released when dropping the `PyClassGuardMap`
138        PyClassGuardMap {
139            ptr: NonNull::from(f(&slf)),
140            checker: slf.as_class_object().borrow_checker(),
141        }
142    }
143}
144
145impl<'a, T> PyClassGuard<'a, T>
146where
147    T: PyClass,
148    T::BaseType: PyClass,
149{
150    /// Borrows a shared reference to `PyClassGuard<T::BaseType>`.
151    ///
152    /// With the help of this method, you can access attributes and call methods
153    /// on the superclass without consuming the `PyClassGuard<T>`. This method
154    /// can also be chained to access the super-superclass (and so on).
155    ///
156    /// # Examples
157    /// ```
158    /// # use pyo3::prelude::*;
159    /// # use pyo3::PyClassGuard;
160    /// #[pyclass(subclass)]
161    /// struct Base {
162    ///     base_name: &'static str,
163    /// }
164    /// #[pymethods]
165    /// impl Base {
166    ///     fn base_name_len(&self) -> usize {
167    ///         self.base_name.len()
168    ///     }
169    /// }
170    ///
171    /// #[pyclass(extends=Base)]
172    /// struct Sub {
173    ///     sub_name: &'static str,
174    /// }
175    ///
176    /// #[pymethods]
177    /// impl Sub {
178    ///     #[new]
179    ///     fn new() -> (Self, Base) {
180    ///         (Self { sub_name: "sub_name" }, Base { base_name: "base_name" })
181    ///     }
182    ///     fn sub_name_len(&self) -> usize {
183    ///         self.sub_name.len()
184    ///     }
185    ///     fn format_name_lengths(slf: PyClassGuard<'_, Self>) -> String {
186    ///         format!("{} {}", slf.as_super().base_name_len(), slf.sub_name_len())
187    ///     }
188    /// }
189    /// # Python::attach(|py| {
190    /// #     let sub = Py::new(py, Sub::new()).unwrap();
191    /// #     pyo3::py_run!(py, sub, "assert sub.format_name_lengths() == '9 8'")
192    /// # });
193    /// ```
194    pub fn as_super(&self) -> &PyClassGuard<'a, T::BaseType> {
195        // SAFETY: `PyClassGuard<T>` and `PyClassGuard<U>` have the same layout
196        unsafe { NonNull::from(self).cast().as_ref() }
197    }
198
199    /// Gets a `PyClassGuard<T::BaseType>`.
200    ///
201    /// With the help of this method, you can get hold of instances of the
202    /// super-superclass when needed.
203    ///
204    /// # Examples
205    /// ```
206    /// # use pyo3::prelude::*;
207    /// # use pyo3::PyClassGuard;
208    /// #[pyclass(subclass)]
209    /// struct Base1 {
210    ///     name1: &'static str,
211    /// }
212    ///
213    /// #[pyclass(extends=Base1, subclass)]
214    /// struct Base2 {
215    ///     name2: &'static str,
216    /// }
217    ///
218    /// #[pyclass(extends=Base2)]
219    /// struct Sub {
220    ///     name3: &'static str,
221    /// }
222    ///
223    /// #[pymethods]
224    /// impl Sub {
225    ///     #[new]
226    ///     fn new() -> PyClassInitializer<Self> {
227    ///         PyClassInitializer::from(Base1 { name1: "base1" })
228    ///             .add_subclass(Base2 { name2: "base2" })
229    ///             .add_subclass(Self { name3: "sub" })
230    ///     }
231    ///     fn name(slf: PyClassGuard<'_, Self>) -> String {
232    ///         let subname = slf.name3;
233    ///         let super_ = slf.into_super();
234    ///         format!("{} {} {}", super_.as_super().name1, super_.name2, subname)
235    ///     }
236    /// }
237    /// # Python::attach(|py| {
238    /// #     let sub = Py::new(py, Sub::new()).unwrap();
239    /// #     pyo3::py_run!(py, sub, "assert sub.name() == 'base1 base2 sub'")
240    /// # });
241    /// ```
242    pub fn into_super(self) -> PyClassGuard<'a, T::BaseType> {
243        let t_not_frozen = !<T::Frozen as crate::pyclass::boolean_struct::private::Boolean>::VALUE;
244        let u_frozen =
245            <<T::BaseType as PyClass>::Frozen as crate::pyclass::boolean_struct::private::Boolean>::VALUE;
246        if t_not_frozen && u_frozen {
247            // If `T` is a mutable subclass of a frozen `U` base, then it is possible that we need
248            // to release the borrow count now. (e.g. `U` may have a noop borrow checker so dropping
249            // the `PyRef<U>` later would noop and leak the borrow we currently hold.)
250            //
251            // However it's nontrivial, if `U` is frozen but itself has a mutable base class `V`,
252            // then the borrow checker of both `T` and `U` is the shared borrow checker of `V`.
253            //
254            // But it's really hard to prove that in the type system, the soundest thing we can do
255            // is just add a borrow to `U` now and then release the borrow of `T`.
256
257            self.as_super()
258                .as_class_object()
259                .borrow_checker()
260                .try_borrow()
261                .expect("this object is already borrowed");
262
263            self.as_class_object().borrow_checker().release_borrow()
264        };
265        PyClassGuard {
266            ptr: std::mem::ManuallyDrop::new(self).ptr,
267            marker: PhantomData,
268        }
269    }
270}
271
272impl<T: PyClass> Deref for PyClassGuard<'_, T> {
273    type Target = T;
274
275    #[inline]
276    fn deref(&self) -> &T {
277        // SAFETY: `PyClassObject<T>` contains a valid `T`, by construction no
278        // mutable alias is enforced
279        unsafe { &*self.as_class_object().get_ptr().cast_const() }
280    }
281}
282
283impl<'a, 'py, T: PyClass> FromPyObject<'a, 'py> for PyClassGuard<'a, T> {
284    type Error = PyClassGuardError<'a, 'py>;
285
286    #[cfg(feature = "experimental-inspect")]
287    const INPUT_TYPE: TypeHint = T::TYPE_HINT;
288
289    fn extract(obj: Borrowed<'a, 'py, crate::PyAny>) -> Result<Self, Self::Error> {
290        Self::try_from_class_object(
291            obj.cast()
292                .map_err(|e| PyClassGuardError(Some(e)))?
293                .get_class_object(),
294        )
295        .map_err(|_| PyClassGuardError(None))
296    }
297}
298
299impl<'a, 'py, T: PyClass> IntoPyObject<'py> for PyClassGuard<'a, T> {
300    type Target = T;
301    type Output = Borrowed<'a, 'py, T>;
302    type Error = Infallible;
303
304    #[cfg(feature = "experimental-inspect")]
305    const OUTPUT_TYPE: TypeHint = T::TYPE_HINT;
306
307    #[inline]
308    fn into_pyobject(self, py: crate::Python<'py>) -> Result<Self::Output, Self::Error> {
309        (&self).into_pyobject(py)
310    }
311}
312
313impl<'a, 'py, T: PyClass> IntoPyObject<'py> for &PyClassGuard<'a, T> {
314    type Target = T;
315    type Output = Borrowed<'a, 'py, T>;
316    type Error = Infallible;
317
318    #[cfg(feature = "experimental-inspect")]
319    const OUTPUT_TYPE: TypeHint = T::TYPE_HINT;
320
321    #[inline]
322    fn into_pyobject(self, py: crate::Python<'py>) -> Result<Self::Output, Self::Error> {
323        // SAFETY: `ptr` is guaranteed to be valid for 'a and points to an
324        // object of type T
325        unsafe { Ok(Borrowed::from_non_null(py, self.ptr).cast_unchecked()) }
326    }
327}
328
329impl<T: PyClass> Drop for PyClassGuard<'_, T> {
330    /// Releases the shared borrow
331    fn drop(&mut self) {
332        self.as_class_object().borrow_checker().release_borrow()
333    }
334}
335
336// SAFETY: `PyClassGuard` only provides access to the inner `T` (and no other
337// Python APIs) which does not require a Python thread state
338#[cfg(feature = "nightly")]
339unsafe impl<T: PyClass> crate::marker::Ungil for PyClassGuard<'_, T> {}
340// SAFETY: we provide access to
341// - `&T`, which requires `T: Sync` to be Send and `T: Sync` to be Sync
342unsafe impl<T: PyClass + Sync> Send for PyClassGuard<'_, T> {}
343unsafe impl<T: PyClass + Sync> Sync for PyClassGuard<'_, T> {}
344
345/// Custom error type for extracting a [PyClassGuard]
346pub struct PyClassGuardError<'a, 'py>(pub(crate) Option<CastError<'a, 'py>>);
347
348impl fmt::Debug for PyClassGuardError<'_, '_> {
349    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
350        if let Some(e) = &self.0 {
351            write!(f, "{e:?}")
352        } else {
353            write!(f, "{:?}", PyBorrowError::new())
354        }
355    }
356}
357
358impl fmt::Display for PyClassGuardError<'_, '_> {
359    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
360        if let Some(e) = &self.0 {
361            write!(f, "{e}")
362        } else {
363            write!(f, "{}", PyBorrowError::new())
364        }
365    }
366}
367
368impl From<PyClassGuardError<'_, '_>> for PyErr {
369    fn from(value: PyClassGuardError<'_, '_>) -> Self {
370        if let Some(e) = value.0 {
371            e.into()
372        } else {
373            PyBorrowError::new().into()
374        }
375    }
376}
377
378/// A wrapper type for a mutably borrowed value from a `PyClass`
379///
380/// # When *not* to use [`PyClassGuardMut`]
381///
382/// Usually you can use `&mut` references as method and function receivers and
383/// arguments, and you won't need to use [`PyClassGuardMut`] directly:
384///
385/// ```rust,no_run
386/// use pyo3::prelude::*;
387///
388/// #[pyclass]
389/// struct Number {
390///     inner: u32,
391/// }
392///
393/// #[pymethods]
394/// impl Number {
395///     fn increment(&mut self) {
396///         self.inner += 1;
397///     }
398/// }
399/// ```
400///
401/// The [`#[pymethods]`](crate::pymethods) proc macro will generate this wrapper
402/// function (and more), using [`PyClassGuardMut`] under the hood:
403///
404/// ```rust,no_run
405/// # use pyo3::prelude::*;
406/// # #[pyclass]
407/// # struct Number {
408/// #    inner: u32,
409/// # }
410/// #
411/// # #[pymethods]
412/// # impl Number {
413/// #    fn increment(&mut self) {
414/// #        self.inner += 1;
415/// #    }
416/// # }
417/// #
418/// // The function which is exported to Python looks roughly like the following
419/// unsafe extern "C" fn __pymethod_increment__(
420///     _slf: *mut ::pyo3::ffi::PyObject,
421///     _args: *mut ::pyo3::ffi::PyObject,
422/// ) -> *mut ::pyo3::ffi::PyObject {
423///     unsafe fn inner<'py>(
424///         py: ::pyo3::Python<'py>,
425///         _slf: *mut ::pyo3::ffi::PyObject,
426///     ) -> ::pyo3::PyResult<*mut ::pyo3::ffi::PyObject> {
427///         let function = Number::increment;
428/// #       #[allow(clippy::let_unit_value)]
429///         let mut holder_0 = ::pyo3::impl_::extract_argument::FunctionArgumentHolder::INIT;
430///         let result = {
431///             let ret = function(::pyo3::impl_::extract_argument::extract_pyclass_ref_mut::<Number>(
432///                 unsafe { ::pyo3::impl_::pymethods::BoundRef::ref_from_ptr(py, &_slf) }.0,
433///                 &mut holder_0,
434///             )?);
435///             {
436///                 let result = {
437///                     let obj = ret;
438/// #                   #[allow(clippy::useless_conversion)]
439///                     ::pyo3::impl_::wrap::converter(&obj)
440///                         .wrap(obj)
441///                         .map_err(::core::convert::Into::<::pyo3::PyErr>::into)
442///                 };
443///                 ::pyo3::impl_::wrap::converter(&result).map_into_ptr(py, result)
444///             }
445///         };
446///         result
447///     }
448///
449///     unsafe {
450///         ::pyo3::impl_::trampoline::get_trampoline_function!(noargs, inner)(
451///             _slf,
452///             _args,
453///         )
454///     }
455/// }
456/// ```
457///
458/// # When to use [`PyClassGuardMut`]
459/// ## Using PyClasses from Rust
460///
461/// However, we *do* need [`PyClassGuardMut`] if we want to call its methods
462/// from Rust:
463/// ```rust
464/// # use pyo3::prelude::*;
465/// # use pyo3::{PyClassGuard, PyClassGuardMut};
466/// #
467/// # #[pyclass]
468/// # struct Number {
469/// #     inner: u32,
470/// # }
471/// #
472/// # #[pymethods]
473/// # impl Number {
474/// #     fn increment(&mut self) {
475/// #         self.inner += 1;
476/// #     }
477/// # }
478/// # fn main() -> PyResult<()> {
479/// Python::attach(|py| {
480///     let n = Py::new(py, Number { inner: 0 })?;
481///
482///     // We borrow the guard and then dereference
483///     // it to get a mutable reference to Number
484///     let mut guard: PyClassGuardMut<'_, Number> = n.extract(py)?;
485///     let n_mutable: &mut Number = &mut *guard;
486///
487///     n_mutable.increment();
488///
489///     // To avoid panics we must dispose of the
490///     // `PyClassGuardMut` before borrowing again.
491///     drop(guard);
492///
493///     let n_immutable: &Number = &*n.extract::<PyClassGuard<'_, Number>>(py)?;
494///     assert_eq!(n_immutable.inner, 1);
495///
496///     Ok(())
497/// })
498/// # }
499/// ```
500/// ## Dealing with possibly overlapping mutable references
501///
502/// It is also necessary to use [`PyClassGuardMut`] if you can receive mutable
503/// arguments that may overlap. Suppose the following function that swaps the
504/// values of two `Number`s:
505/// ```
506/// # use pyo3::prelude::*;
507/// # #[pyclass]
508/// # pub struct Number {
509/// #     inner: u32,
510/// # }
511/// #[pyfunction]
512/// fn swap_numbers(a: &mut Number, b: &mut Number) {
513///     std::mem::swap(&mut a.inner, &mut b.inner);
514/// }
515/// # fn main() {
516/// #     Python::attach(|py| {
517/// #         let n = Py::new(py, Number{inner: 35}).unwrap();
518/// #         let n2 = n.clone_ref(py);
519/// #         assert!(n.is(&n2));
520/// #         let fun = pyo3::wrap_pyfunction!(swap_numbers, py).unwrap();
521/// #         fun.call1((n, n2)).expect_err("Managed to create overlapping mutable references. Note: this is undefined behaviour.");
522/// #     });
523/// # }
524/// ```
525/// When users pass in the same `Number` as both arguments, one of the mutable
526/// borrows will fail and raise a `RuntimeError`:
527/// ```text
528/// >>> a = Number()
529/// >>> swap_numbers(a, a)
530/// Traceback (most recent call last):
531///   File "<stdin>", line 1, in <module>
532///   RuntimeError: Already borrowed
533/// ```
534///
535/// It is better to write that function like this:
536/// ```rust
537/// # use pyo3::prelude::*;
538/// # use pyo3::{PyClassGuard, PyClassGuardMut};
539/// # #[pyclass]
540/// # pub struct Number {
541/// #     inner: u32,
542/// # }
543/// #[pyfunction]
544/// fn swap_numbers(a: &Bound<'_, Number>, b: &Bound<'_, Number>) -> PyResult<()> {
545///     // Check that the pointers are unequal
546///     if !a.is(b) {
547///         let mut a: PyClassGuardMut<'_, Number> = a.extract()?;
548///         let mut b: PyClassGuardMut<'_, Number> = b.extract()?;
549///         std::mem::swap(&mut a.inner, &mut b.inner);
550///     } else {
551///         // Do nothing - they are the same object, so don't need swapping.
552///     }
553///     Ok(())
554/// }
555/// # fn main() {
556/// #     // With duplicate numbers
557/// #     Python::attach(|py| {
558/// #         let n = Py::new(py, Number{inner: 35}).unwrap();
559/// #         let n2 = n.clone_ref(py);
560/// #         assert!(n.is(&n2));
561/// #         let fun = pyo3::wrap_pyfunction!(swap_numbers, py).unwrap();
562/// #         fun.call1((n, n2)).unwrap();
563/// #     });
564/// #
565/// #     // With two different numbers
566/// #     Python::attach(|py| {
567/// #         let n = Py::new(py, Number{inner: 35}).unwrap();
568/// #         let n2 = Py::new(py, Number{inner: 42}).unwrap();
569/// #         assert!(!n.is(&n2));
570/// #         let fun = pyo3::wrap_pyfunction!(swap_numbers, py).unwrap();
571/// #         fun.call1((&n, &n2)).unwrap();
572/// #         let n: u32 = n.extract::<PyClassGuard<'_, Number>>(py).unwrap().inner;
573/// #         let n2: u32 = n2.extract::<PyClassGuard<'_, Number>>(py).unwrap().inner;
574/// #         assert_eq!(n, 42);
575/// #         assert_eq!(n2, 35);
576/// #     });
577/// # }
578/// ```
579/// See [`PyClassGuard`] and the [guide] for more information.
580///
581/// [guide]: https://pyo3.rs/latest/class.html#bound-and-interior-mutability
582///     "Bound and interior mutability"
583#[repr(transparent)]
584pub struct PyClassGuardMut<'a, T: PyClass<Frozen = False>> {
585    ptr: NonNull<ffi::PyObject>,
586    marker: PhantomData<&'a Py<T>>,
587}
588
589impl<'a, T: PyClass<Frozen = False>> PyClassGuardMut<'a, T> {
590    pub(crate) fn try_borrow_mut(obj: &'a Py<T>) -> Result<Self, PyBorrowMutError> {
591        Self::try_from_class_object(obj.get_class_object())
592    }
593
594    fn try_from_class_object(obj: &'a PyClassObject<T>) -> Result<Self, PyBorrowMutError> {
595        obj.ensure_threadsafe();
596        obj.borrow_checker().try_borrow_mut().map(|_| Self {
597            ptr: NonNull::from(obj).cast(),
598            marker: PhantomData,
599        })
600    }
601
602    pub(crate) fn as_class_object(&self) -> &'a PyClassObject<T> {
603        // SAFETY: `ptr` by construction points to a `PyClassObject<T>` and is
604        // valid for at least 'a
605        unsafe { self.ptr.cast().as_ref() }
606    }
607
608    /// Consumes the [`PyClassGuardMut`] and returns a [`PyClassGuardMap`] for a component of the
609    /// borrowed data
610    ///
611    /// # Examples
612    ///
613    /// ```
614    /// # use pyo3::prelude::*;
615    /// # use pyo3::PyClassGuardMut;
616    ///
617    /// #[pyclass]
618    /// pub struct MyClass {
619    ///     data: [i32; 100],
620    /// }
621    ///
622    /// # Python::attach(|py| {
623    /// let obj = Bound::new(py, MyClass { data: [0; 100] })?;
624    /// let mut data = obj.extract::<PyClassGuardMut<'_, MyClass>>()?.map(|c| c.data.as_mut_slice());
625    /// data[0] = 42;
626    /// # Ok::<_, PyErr>(())
627    /// # }).unwrap();
628    /// ```
629    pub fn map<F, U: ?Sized>(self, f: F) -> PyClassGuardMap<'a, U, true>
630    where
631        F: FnOnce(&mut T) -> &mut U,
632    {
633        let mut slf = std::mem::ManuallyDrop::new(self); // the borrow is released when dropping the `PyClassGuardMap`
634        PyClassGuardMap {
635            ptr: NonNull::from(f(&mut slf)),
636            checker: slf.as_class_object().borrow_checker(),
637        }
638    }
639}
640
641impl<'a, T> PyClassGuardMut<'a, T>
642where
643    T: PyClass<Frozen = False>,
644    T::BaseType: PyClass<Frozen = False>,
645{
646    /// Borrows a mutable reference to `PyClassGuardMut<T::BaseType>`.
647    ///
648    /// With the help of this method, you can mutate attributes and call
649    /// mutating methods on the superclass without consuming the
650    /// `PyClassGuardMut<T>`. This method can also be chained to access the
651    /// super-superclass (and so on).
652    ///
653    /// See [`PyClassGuard::as_super`] for more.
654    pub fn as_super(&mut self) -> &mut PyClassGuardMut<'a, T::BaseType> {
655        // SAFETY: `PyClassGuardMut<T>` and `PyClassGuardMut<U>` have the same layout
656        unsafe { NonNull::from(self).cast().as_mut() }
657    }
658
659    /// Gets a `PyClassGuardMut<T::BaseType>`.
660    ///
661    /// See [`PyClassGuard::into_super`] for more.
662    pub fn into_super(self) -> PyClassGuardMut<'a, T::BaseType> {
663        // `PyClassGuardMut` is only available for non-frozen classes, so there
664        // is no possibility of leaking borrows like `PyClassGuard`
665        PyClassGuardMut {
666            ptr: std::mem::ManuallyDrop::new(self).ptr,
667            marker: PhantomData,
668        }
669    }
670}
671
672impl<T: PyClass<Frozen = False>> Deref for PyClassGuardMut<'_, T> {
673    type Target = T;
674
675    #[inline]
676    fn deref(&self) -> &T {
677        // SAFETY: `PyClassObject<T>` contains a valid `T`, by construction no
678        // alias is enforced
679        unsafe { &*self.as_class_object().get_ptr().cast_const() }
680    }
681}
682impl<T: PyClass<Frozen = False>> DerefMut for PyClassGuardMut<'_, T> {
683    #[inline]
684    fn deref_mut(&mut self) -> &mut T {
685        // SAFETY: `PyClassObject<T>` contains a valid `T`, by construction no
686        // alias is enforced
687        unsafe { &mut *self.as_class_object().get_ptr() }
688    }
689}
690
691impl<'a, 'py, T: PyClass<Frozen = False>> FromPyObject<'a, 'py> for PyClassGuardMut<'a, T> {
692    type Error = PyClassGuardMutError<'a, 'py>;
693
694    #[cfg(feature = "experimental-inspect")]
695    const INPUT_TYPE: TypeHint = T::TYPE_HINT;
696
697    fn extract(obj: Borrowed<'a, 'py, crate::PyAny>) -> Result<Self, Self::Error> {
698        Self::try_from_class_object(
699            obj.cast()
700                .map_err(|e| PyClassGuardMutError(Some(e)))?
701                .get_class_object(),
702        )
703        .map_err(|_| PyClassGuardMutError(None))
704    }
705}
706
707impl<'a, 'py, T: PyClass<Frozen = False>> IntoPyObject<'py> for PyClassGuardMut<'a, T> {
708    type Target = T;
709    type Output = Borrowed<'a, 'py, T>;
710    type Error = Infallible;
711
712    #[cfg(feature = "experimental-inspect")]
713    const OUTPUT_TYPE: TypeHint = T::TYPE_HINT;
714
715    #[inline]
716    fn into_pyobject(self, py: crate::Python<'py>) -> Result<Self::Output, Self::Error> {
717        (&self).into_pyobject(py)
718    }
719}
720
721impl<'a, 'py, T: PyClass<Frozen = False>> IntoPyObject<'py> for &PyClassGuardMut<'a, T> {
722    type Target = T;
723    type Output = Borrowed<'a, 'py, T>;
724    type Error = Infallible;
725
726    #[cfg(feature = "experimental-inspect")]
727    const OUTPUT_TYPE: TypeHint = T::TYPE_HINT;
728
729    #[inline]
730    fn into_pyobject(self, py: crate::Python<'py>) -> Result<Self::Output, Self::Error> {
731        // SAFETY: `ptr` is guaranteed to be valid for 'a and points to an
732        // object of type T
733        unsafe { Ok(Borrowed::from_non_null(py, self.ptr).cast_unchecked()) }
734    }
735}
736
737impl<T: PyClass<Frozen = False>> Drop for PyClassGuardMut<'_, T> {
738    /// Releases the mutable borrow
739    fn drop(&mut self) {
740        self.as_class_object().borrow_checker().release_borrow_mut()
741    }
742}
743
744// SAFETY: `PyClassGuardMut` only provides access to the inner `T` (and no other
745// Python APIs) which does not require a Python thread state
746#[cfg(feature = "nightly")]
747unsafe impl<T: PyClass<Frozen = False>> crate::marker::Ungil for PyClassGuardMut<'_, T> {}
748// SAFETY: we provide access to
749// - `&T`, which requires `T: Sync` to be Send and `T: Sync` to be Sync
750// - `&mut T`, which requires `T: Send` to be Send and `T: Sync` to be Sync
751unsafe impl<T: PyClass<Frozen = False> + Send + Sync> Send for PyClassGuardMut<'_, T> {}
752unsafe impl<T: PyClass<Frozen = False> + Sync> Sync for PyClassGuardMut<'_, T> {}
753
754/// Custom error type for extracting a [PyClassGuardMut]
755pub struct PyClassGuardMutError<'a, 'py>(pub(crate) Option<CastError<'a, 'py>>);
756
757impl fmt::Debug for PyClassGuardMutError<'_, '_> {
758    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
759        if let Some(e) = &self.0 {
760            write!(f, "{e:?}")
761        } else {
762            write!(f, "{:?}", PyBorrowMutError::new())
763        }
764    }
765}
766
767impl fmt::Display for PyClassGuardMutError<'_, '_> {
768    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
769        if let Some(e) = &self.0 {
770            write!(f, "{e}")
771        } else {
772            write!(f, "{}", PyBorrowMutError::new())
773        }
774    }
775}
776
777impl From<PyClassGuardMutError<'_, '_>> for PyErr {
778    fn from(value: PyClassGuardMutError<'_, '_>) -> Self {
779        if let Some(e) = value.0 {
780            e.into()
781        } else {
782            PyBorrowMutError::new().into()
783        }
784    }
785}
786
787/// Wraps a borrowed reference `U` to a value stored inside of a pyclass `T`
788///
789/// See [`PyClassGuard::map`] and [`PyClassGuardMut::map`]
790pub struct PyClassGuardMap<'a, U: ?Sized, const MUT: bool> {
791    ptr: NonNull<U>,
792    checker: &'a dyn PyClassBorrowChecker,
793}
794
795impl<U: ?Sized, const MUT: bool> Deref for PyClassGuardMap<'_, U, MUT> {
796    type Target = U;
797
798    fn deref(&self) -> &U {
799        // SAFETY: `checker` guards our access to the `T` that `U` points into
800        unsafe { self.ptr.as_ref() }
801    }
802}
803
804impl<U: ?Sized> DerefMut for PyClassGuardMap<'_, U, true> {
805    fn deref_mut(&mut self) -> &mut Self::Target {
806        // SAFETY: `checker` guards our access to the `T` that `U` points into
807        unsafe { self.ptr.as_mut() }
808    }
809}
810
811impl<U: ?Sized, const MUT: bool> Drop for PyClassGuardMap<'_, U, MUT> {
812    fn drop(&mut self) {
813        if MUT {
814            self.checker.release_borrow_mut();
815        } else {
816            self.checker.release_borrow();
817        }
818    }
819}
820
821#[cfg(test)]
822#[cfg(feature = "macros")]
823mod tests {
824    use super::{PyClassGuard, PyClassGuardMut};
825    use crate::{types::PyAnyMethods as _, Bound, IntoPyObject as _, Py, PyErr, Python};
826
827    #[test]
828    fn test_into_frozen_super_released_borrow() {
829        #[crate::pyclass]
830        #[pyo3(crate = "crate", subclass, frozen)]
831        struct BaseClass {}
832
833        #[crate::pyclass]
834        #[pyo3(crate = "crate", extends=BaseClass, subclass)]
835        struct SubClass {}
836
837        #[crate::pymethods]
838        #[pyo3(crate = "crate")]
839        impl SubClass {
840            #[new]
841            fn new(py: Python<'_>) -> Py<SubClass> {
842                let init = crate::PyClassInitializer::from(BaseClass {}).add_subclass(SubClass {});
843                Py::new(py, init).expect("allocation error")
844            }
845        }
846
847        Python::attach(|py| {
848            let obj = SubClass::new(py);
849            drop(PyClassGuard::try_borrow(&obj).unwrap().into_super());
850            assert!(PyClassGuardMut::try_borrow_mut(&obj).is_ok());
851        })
852    }
853
854    #[test]
855    fn test_into_frozen_super_mutable_base_holds_borrow() {
856        #[crate::pyclass]
857        #[pyo3(crate = "crate", subclass)]
858        struct BaseClass {}
859
860        #[crate::pyclass]
861        #[pyo3(crate = "crate", extends=BaseClass, subclass, frozen)]
862        struct SubClass {}
863
864        #[crate::pyclass]
865        #[pyo3(crate = "crate", extends=SubClass, subclass)]
866        struct SubSubClass {}
867
868        #[crate::pymethods]
869        #[pyo3(crate = "crate")]
870        impl SubSubClass {
871            #[new]
872            fn new(py: Python<'_>) -> Py<SubSubClass> {
873                let init = crate::PyClassInitializer::from(BaseClass {})
874                    .add_subclass(SubClass {})
875                    .add_subclass(SubSubClass {});
876                Py::new(py, init).expect("allocation error")
877            }
878        }
879
880        Python::attach(|py| {
881            let obj = SubSubClass::new(py);
882            let _super_borrow = PyClassGuard::try_borrow(&obj).unwrap().into_super();
883            // the whole object still has an immutable borrow, so we cannot
884            // borrow any part mutably (the borrowflag is shared)
885            assert!(PyClassGuardMut::try_borrow_mut(&obj).is_err());
886        })
887    }
888
889    #[crate::pyclass]
890    #[pyo3(crate = "crate", subclass)]
891    struct BaseClass {
892        val1: usize,
893    }
894
895    #[crate::pyclass]
896    #[pyo3(crate = "crate", extends=BaseClass, subclass)]
897    struct SubClass {
898        val2: usize,
899    }
900
901    #[crate::pyclass]
902    #[pyo3(crate = "crate", extends=SubClass)]
903    struct SubSubClass {
904        #[pyo3(get)]
905        val3: usize,
906    }
907
908    #[crate::pymethods]
909    #[pyo3(crate = "crate")]
910    impl SubSubClass {
911        #[new]
912        fn new(py: Python<'_>) -> Py<SubSubClass> {
913            let init = crate::PyClassInitializer::from(BaseClass { val1: 10 })
914                .add_subclass(SubClass { val2: 15 })
915                .add_subclass(SubSubClass { val3: 20 });
916            Py::new(py, init).expect("allocation error")
917        }
918
919        fn get_values(self_: PyClassGuard<'_, Self>) -> (usize, usize, usize) {
920            let val1 = self_.as_super().as_super().val1;
921            let val2 = self_.as_super().val2;
922            (val1, val2, self_.val3)
923        }
924
925        fn double_values(mut self_: PyClassGuardMut<'_, Self>) {
926            self_.as_super().as_super().val1 *= 2;
927            self_.as_super().val2 *= 2;
928            self_.val3 *= 2;
929        }
930
931        fn __add__<'a>(
932            mut slf: PyClassGuardMut<'a, Self>,
933            other: PyClassGuard<'a, Self>,
934        ) -> PyClassGuardMut<'a, Self> {
935            slf.val3 += other.val3;
936            slf
937        }
938
939        fn __rsub__<'a>(
940            slf: PyClassGuard<'a, Self>,
941            mut other: PyClassGuardMut<'a, Self>,
942        ) -> PyClassGuardMut<'a, Self> {
943            other.val3 -= slf.val3;
944            other
945        }
946    }
947
948    #[test]
949    fn test_pyclassguard_into_pyobject() {
950        Python::attach(|py| {
951            let class = Py::new(py, BaseClass { val1: 42 })?;
952            let guard = PyClassGuard::try_borrow(&class).unwrap();
953            let new_ref = (&guard).into_pyobject(py)?;
954            assert!(new_ref.is(&class));
955            let new = guard.into_pyobject(py)?;
956            assert!(new.is(&class));
957            Ok::<_, PyErr>(())
958        })
959        .unwrap();
960    }
961
962    #[test]
963    fn test_pyclassguardmut_into_pyobject() {
964        Python::attach(|py| {
965            let class = Py::new(py, BaseClass { val1: 42 })?;
966            let guard = PyClassGuardMut::try_borrow_mut(&class).unwrap();
967            let new_ref = (&guard).into_pyobject(py)?;
968            assert!(new_ref.is(&class));
969            let new = guard.into_pyobject(py)?;
970            assert!(new.is(&class));
971            Ok::<_, PyErr>(())
972        })
973        .unwrap();
974    }
975    #[test]
976    fn test_pyclassguard_as_super() {
977        Python::attach(|py| {
978            let obj = SubSubClass::new(py).into_bound(py);
979            let pyref = PyClassGuard::try_borrow(obj.as_unbound()).unwrap();
980            assert_eq!(pyref.as_super().as_super().val1, 10);
981            assert_eq!(pyref.as_super().val2, 15);
982            assert_eq!(pyref.val3, 20);
983            assert_eq!(SubSubClass::get_values(pyref), (10, 15, 20));
984        });
985    }
986
987    #[test]
988    fn test_pyclassguardmut_as_super() {
989        Python::attach(|py| {
990            let obj = SubSubClass::new(py).into_bound(py);
991            assert_eq!(
992                SubSubClass::get_values(PyClassGuard::try_borrow(obj.as_unbound()).unwrap()),
993                (10, 15, 20)
994            );
995            {
996                let mut pyrefmut = PyClassGuardMut::try_borrow_mut(obj.as_unbound()).unwrap();
997                assert_eq!(pyrefmut.as_super().as_super().val1, 10);
998                pyrefmut.as_super().as_super().val1 -= 5;
999                pyrefmut.as_super().val2 -= 5;
1000                pyrefmut.val3 -= 5;
1001            }
1002            assert_eq!(
1003                SubSubClass::get_values(PyClassGuard::try_borrow(obj.as_unbound()).unwrap()),
1004                (5, 10, 15)
1005            );
1006            SubSubClass::double_values(PyClassGuardMut::try_borrow_mut(obj.as_unbound()).unwrap());
1007            assert_eq!(
1008                SubSubClass::get_values(PyClassGuard::try_borrow(obj.as_unbound()).unwrap()),
1009                (10, 20, 30)
1010            );
1011        });
1012    }
1013
1014    #[test]
1015    fn test_extract_guard() {
1016        Python::attach(|py| {
1017            let obj1 = SubSubClass::new(py);
1018            let obj2 = SubSubClass::new(py);
1019            crate::py_run!(py, obj1 obj2, "assert ((obj1 + obj2) - obj2).val3 == obj1.val3");
1020        });
1021    }
1022
1023    #[test]
1024    fn test_pyclassguards_in_python() {
1025        Python::attach(|py| {
1026            let obj = SubSubClass::new(py);
1027            crate::py_run!(py, obj, "assert obj.get_values() == (10, 15, 20)");
1028            crate::py_run!(py, obj, "assert obj.double_values() is None");
1029            crate::py_run!(py, obj, "assert obj.get_values() == (20, 30, 40)");
1030        });
1031    }
1032
1033    #[crate::pyclass]
1034    #[pyo3(crate = "crate")]
1035    pub struct MyClass {
1036        data: [i32; 100],
1037    }
1038
1039    #[test]
1040    fn test_pyclassguard_map() {
1041        Python::attach(|py| {
1042            let obj = Bound::new(py, MyClass { data: [0; 100] })?;
1043            let data = PyClassGuard::try_borrow(obj.as_unbound())?.map(|c| &c.data);
1044            assert_eq!(data[0], 0);
1045            assert!(obj.try_borrow_mut().is_err()); // obj is still protected
1046            drop(data);
1047            assert!(obj.try_borrow_mut().is_ok()); // drop released shared borrow
1048            Ok::<_, PyErr>(())
1049        })
1050        .unwrap()
1051    }
1052
1053    #[test]
1054    fn test_pyclassguardmut_map() {
1055        Python::attach(|py| {
1056            let obj = Bound::new(py, MyClass { data: [0; 100] })?;
1057            let mut data =
1058                PyClassGuardMut::try_borrow_mut(obj.as_unbound())?.map(|c| c.data.as_mut_slice());
1059            assert_eq!(data[0], 0);
1060            data[0] = 5;
1061            assert_eq!(data[0], 5);
1062            assert!(obj.try_borrow_mut().is_err()); // obj is still protected
1063            drop(data);
1064            assert!(obj.try_borrow_mut().is_ok()); // drop released mutable borrow
1065            Ok::<_, PyErr>(())
1066        })
1067        .unwrap()
1068    }
1069
1070    #[test]
1071    fn test_pyclassguard_map_unrelated() {
1072        use crate::types::{PyString, PyStringMethods};
1073        Python::attach(|py| {
1074            let obj = Bound::new(py, MyClass { data: [0; 100] })?;
1075            let string = PyString::new(py, "pyo3");
1076            // It is possible to return something not borrowing from the guard, but that shouldn't
1077            // matter. `RefCell` has the same behaviour
1078            let refmap = PyClassGuard::try_borrow(obj.as_unbound())?.map(|_| &string);
1079            assert_eq!(refmap.to_cow()?, "pyo3");
1080            Ok::<_, PyErr>(())
1081        })
1082        .unwrap()
1083    }
1084}