pyo3/pyclass/
guard.rs

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