Skip to main content

pyo3/pyclass/
guard.rs

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