pyo3_utils/
py_wrapper.rs

1//! When you want to implement a pyclass that can obtain ownership of the internal value in a semantically compatible way,
2//! you can use this module.
3//!
4//! Pay attention to this module's:
5//!
6//! - [PyWrapper]
7//! - [PyWrapperSemverExt]
8
9use std::convert::Infallible;
10use std::error::Error;
11use std::fmt::{Debug, Display, Formatter};
12use std::mem::replace;
13use std::ops::{Deref, DerefMut};
14
15use parking_lot::{
16    MappedRwLockReadGuard, MappedRwLockWriteGuard, RwLock, RwLockReadGuard, RwLockWriteGuard,
17};
18use pyo3::exceptions::PyRuntimeError;
19use pyo3::PyErr;
20
21const CONSUMED_ERROR_MSG: &str = "Already consumed";
22const LOCK_ERROR_MSG: &str = "Already mutably borrowed";
23
24/// This error indicates that the internal value has been consumed, i.e., its ownership has been moved out.
25#[derive(Debug)]
26pub struct ConsumedError;
27
28impl Display for ConsumedError {
29    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
30        write!(f, "{CONSUMED_ERROR_MSG}")
31    }
32}
33
34impl Error for ConsumedError {}
35
36impl From<ConsumedError> for PyErr {
37    fn from(_: ConsumedError) -> Self {
38        PyRuntimeError::new_err(CONSUMED_ERROR_MSG)
39    }
40}
41
42pub type ConsumedResult<T> = Result<T, ConsumedError>;
43
44/// This error indicates that the internal value has already been mutably borrowed.
45/// At this point, you must wait for other code to release the lock.
46#[derive(Debug)]
47pub struct LockError;
48
49impl Display for LockError {
50    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
51        write!(f, "{LOCK_ERROR_MSG}")
52    }
53}
54
55impl Error for LockError {}
56
57impl From<LockError> for PyErr {
58    fn from(_: LockError) -> Self {
59        PyRuntimeError::new_err(LOCK_ERROR_MSG)
60    }
61}
62
63pub type LockResult<T> = Result<T, LockError>;
64
65/// Can only obtain alias references
66pub type PyWrapperT0<T> = Result<T, Infallible>;
67/// Can obtain alias references and mutable references
68// TODO, FIXME, PERF: we have to use `RwLock` instead of `RefCell`,
69// it's because pyo3 require `Sync`.
70//
71// We need wait for pyo3 `unsync`, see also:
72// - <https://github.com/PyO3/pyo3/issues/4265#issuecomment-2348510879>
73// - <https://github.com/pydantic/pydantic-core/pull/1556#issue-2694035224>
74//
75// ---
76//
77// use `parking_lot` instead of `std`, it's because `parking_lot` will not poisoned.
78pub type PyWrapperT1<T> = RwLock<Result<T, Infallible>>;
79/// Can obtain alias references, mutable references, and ownership
80pub type PyWrapperT2<T> = RwLock<Result<T, ConsumedError>>;
81
82mod sealed {
83    use super::*;
84
85    pub trait PyWrapperT {}
86
87    impl<T> PyWrapperT for PyWrapperT0<T> {}
88    impl<T> PyWrapperT for PyWrapperT1<T> {}
89    impl<T> PyWrapperT for PyWrapperT2<T> {}
90
91    pub trait SealedPyWrapper {}
92
93    impl<T> SealedPyWrapper for PyWrapper<T> where T: PyWrapperT {}
94
95    pub trait SealedMappableDeref {}
96
97    impl<T: ?Sized> SealedMappableDeref for &T {}
98    impl<T: ?Sized> SealedMappableDeref for RwLockReadGuard<'_, T> {}
99    impl<T: ?Sized> SealedMappableDeref for MappedRwLockReadGuard<'_, T> {}
100
101    pub trait SealedMappableDerefMut {}
102
103    impl<T: ?Sized> SealedMappableDerefMut for &mut T {}
104    impl<T: ?Sized> SealedMappableDerefMut for RwLockWriteGuard<'_, T> {}
105    impl<T: ?Sized> SealedMappableDerefMut for MappedRwLockWriteGuard<'_, T> {}
106}
107
108trait RwLockExt {
109    type T;
110
111    fn try_read_ext(&self) -> LockResult<RwLockReadGuard<'_, Self::T>>;
112
113    fn try_write_ext(&self) -> LockResult<RwLockWriteGuard<'_, Self::T>>;
114}
115
116impl<T> RwLockExt for RwLock<T> {
117    type T = T;
118
119    fn try_read_ext(&self) -> LockResult<RwLockReadGuard<'_, T>> {
120        self.try_read().ok_or(LockError)
121    }
122
123    fn try_write_ext(&self) -> LockResult<RwLockWriteGuard<'_, T>> {
124        self.try_write().ok_or(LockError)
125    }
126}
127
128/// This trait provides compatibility between `&T` and [parking_lot::RwLockReadGuard]
129pub trait MappableDeref<'a>: Deref + sealed::SealedMappableDeref {
130    /// This method is similar to [parking_lot::RwLockReadGuard::map] and its sibling methods.
131    fn map<U, F>(self, f: F) -> impl MappableDeref<'a, Target = U>
132    where
133        U: ?Sized + 'a,
134        F: FnOnce(&Self::Target) -> &U;
135}
136
137impl<'a, T> MappableDeref<'a> for &'a T
138where
139    T: ?Sized,
140{
141    fn map<U, F>(self, f: F) -> impl MappableDeref<'a, Target = U>
142    where
143        U: ?Sized + 'a,
144        F: FnOnce(&T) -> &U,
145    {
146        f(self)
147    }
148}
149
150impl<'a, T> MappableDeref<'a> for MappedRwLockReadGuard<'a, T>
151where
152    T: ?Sized + 'a,
153{
154    fn map<U, F>(self, f: F) -> impl MappableDeref<'a, Target = U>
155    where
156        U: ?Sized + 'a,
157        F: FnOnce(&T) -> &U,
158    {
159        MappedRwLockReadGuard::map(self, f)
160    }
161}
162
163impl<'a, T> MappableDeref<'a> for RwLockReadGuard<'a, T>
164where
165    T: ?Sized + 'a,
166{
167    fn map<U, F>(self, f: F) -> impl MappableDeref<'a, Target = U>
168    where
169        U: ?Sized + 'a,
170        F: FnOnce(&T) -> &U,
171    {
172        RwLockReadGuard::map(self, f)
173    }
174}
175
176/// This trait provides compatibility between [&mut T] and [parking_lot::RwLockWriteGuard]
177pub trait MappableDerefMut<'a>: DerefMut + sealed::SealedMappableDerefMut {
178    /// This method is similar to [parking_lot::RwLockWriteGuard::map] and its sibling methods.
179    fn map<U, F>(self, f: F) -> impl MappableDerefMut<'a, Target = U>
180    where
181        U: ?Sized + 'a,
182        F: FnOnce(&mut Self::Target) -> &mut U;
183}
184
185impl<'a, T> MappableDerefMut<'a> for &'a mut T
186where
187    T: ?Sized,
188{
189    fn map<U, F>(self, f: F) -> impl MappableDerefMut<'a, Target = U>
190    where
191        U: ?Sized + 'a,
192        F: FnOnce(&mut T) -> &mut U,
193    {
194        f(self)
195    }
196}
197
198impl<'a, T> MappableDerefMut<'a> for MappedRwLockWriteGuard<'a, T>
199where
200    T: ?Sized + 'a,
201{
202    fn map<U, F>(self, f: F) -> impl MappableDerefMut<'a, Target = U>
203    where
204        U: ?Sized + 'a,
205        F: FnOnce(&mut T) -> &mut U,
206    {
207        MappedRwLockWriteGuard::map(self, f)
208    }
209}
210
211impl<'a, T> MappableDerefMut<'a> for RwLockWriteGuard<'a, T>
212where
213    T: ?Sized + 'a,
214{
215    fn map<U, F>(self, f: F) -> impl MappableDerefMut<'a, Target = U>
216    where
217        U: ?Sized + 'a,
218        F: FnOnce(&mut T) -> &mut U,
219    {
220        RwLockWriteGuard::map(self, f)
221    }
222}
223
224/// You can wrap the desired internal value in this structure to implement a pyclass that
225/// can obtain ownership of the internal value.
226///
227/// # Example
228/**
229```rust
230use pyo3::prelude::*;
231use pyo3_utils::py_wrapper::{PyWrapper, PyWrapperT2};
232
233struct Foo;
234
235impl Foo {
236    fn foo(self) {}
237}
238
239#[pyclass(frozen)]
240#[non_exhaustive]
241pub struct Bar(PyWrapper<PyWrapperT2<Foo>>);
242
243#[pymethods]
244impl Bar {
245    // Normally you can directly use `&self` to get a reference
246    // instead of using `Py<Self>::get` in a pymethod.
247    // Here just for demonstration purposes.
248    fn py_foo(slf: Py<Self>) -> PyResult<()> {
249        slf.get().0.try_take_inner()??.foo();
250        Ok(())
251    }
252}
253```
254*/
255/// NOTE: For [`PyWrapper<T>`], changes from `T = [PyWrapperT0] -> [PyWrapperT1] -> [PyWrapperT2]`
256/// will not be considered breaking changes.
257///
258/// - When the type is [PyWrapperT0], all methods are zero-cost abstractions.
259/// - When the type changes to [PyWrapperT1], compatibility with [PyWrapperT0] is achieved by
260///   implicitly calling other methods that acquire locks. These compatible methods will emit
261///   deprecation warnings.
262/// - When the type changes to [PyWrapperT2], compatibility with [PyWrapperT1] is achieved by
263///   implicitly calling [Result::unwrap()] on other methods that return [Result]. These
264///   compatible methods will emit deprecation warnings.
265pub struct PyWrapper<T>
266where
267    T: sealed::PyWrapperT,
268{
269    inner: T,
270}
271
272impl<T> PyWrapper<PyWrapperT0<T>> {
273    #[inline]
274    pub fn new0(inner: T) -> Self {
275        Self { inner: Ok(inner) }
276    }
277
278    #[inline]
279    pub fn inner_ref(&self) -> impl MappableDeref<'_, Target = T> {
280        // TODO, FIXME: use [Result::into_ok] instead (unstable for now)
281        self.inner.as_ref().unwrap()
282    }
283
284    #[inline]
285    pub fn inner_mut(&mut self) -> impl MappableDerefMut<'_, Target = T> {
286        // TODO, FIXME: use [Result::into_ok] instead (unstable for now)
287        self.inner.as_mut().unwrap()
288    }
289
290    #[inline]
291    pub fn into_inner(self) -> T {
292        // TODO, FIXME: use [Result::into_ok] instead (unstable for now)
293        self.inner.unwrap()
294    }
295}
296
297impl<T> PyWrapper<PyWrapperT1<T>> {
298    #[inline]
299    pub fn new1(inner: T) -> Self {
300        Self {
301            inner: RwLock::new(Ok(inner)),
302        }
303    }
304
305    pub fn lock_inner_ref(&self) -> LockResult<MappedRwLockReadGuard<'_, T>> {
306        self.inner
307            .try_read_ext()
308            // TODO, FIXME: use [Result::into_ok] instead (unstable for now)
309            .map(|guard| RwLockReadGuard::map(guard, |inner| inner.as_ref().unwrap()))
310    }
311
312    pub fn lock_inner_mut(&self) -> LockResult<MappedRwLockWriteGuard<'_, T>> {
313        self.inner
314            .try_write_ext()
315            // TODO, FIXME: use [Result::into_ok] instead (unstable for now)
316            .map(|guard| RwLockWriteGuard::map(guard, |inner| inner.as_mut().unwrap()))
317    }
318
319    pub fn into_inner(self) -> T {
320        // TODO, FIXME: use [Result::into_ok] instead (unstable for now)
321        self.inner.into_inner().unwrap()
322    }
323
324    /// # Panics
325    ///
326    /// Panics if the internal value has already been mutably borrowed.
327    #[deprecated(note = "use `lock_inner_ref` instead")]
328    pub fn inner_ref(&self) -> impl MappableDeref<'_, Target = T> {
329        self.lock_inner_ref().expect(LOCK_ERROR_MSG)
330    }
331
332    /// # Panics
333    ///
334    /// Panics if the internal value has already been mutably borrowed.
335    #[deprecated(note = "use `lock_inner_mut` instead")]
336    pub fn inner_mut(&self) -> impl MappableDerefMut<'_, Target = T> {
337        self.lock_inner_mut().expect(LOCK_ERROR_MSG)
338    }
339}
340
341impl<T> PyWrapper<PyWrapperT2<T>> {
342    #[inline]
343    pub fn new2(inner: T) -> Self {
344        Self {
345            inner: RwLock::new(Ok(inner)),
346        }
347    }
348
349    pub fn try_lock_inner_ref(&self) -> LockResult<ConsumedResult<MappedRwLockReadGuard<'_, T>>> {
350        self.try_read().map(|guard| {
351            if guard.is_err() {
352                Err(ConsumedError)
353            } else {
354                // PEFR: it's ok to use [unwrap_unchecked], but i dont like unsafe block
355                Ok(RwLockReadGuard::map(guard, |inner| inner.as_ref().unwrap()))
356            }
357        })
358    }
359
360    pub fn try_lock_inner_mut(&self) -> LockResult<ConsumedResult<MappedRwLockWriteGuard<'_, T>>> {
361        self.try_write().map(|guard| {
362            if guard.is_err() {
363                Err(ConsumedError)
364            } else {
365                // PEFR: it's ok to use [unwrap_unchecked], but i dont like unsafe block
366                Ok(RwLockWriteGuard::map(guard, |inner| {
367                    inner.as_mut().unwrap()
368                }))
369            }
370        })
371    }
372
373    pub fn try_take_inner(&self) -> LockResult<ConsumedResult<T>> {
374        self.try_replace_inner(Err(ConsumedError))
375    }
376
377    /// similar to [std::mem::replace]
378    pub fn try_replace_inner(&self, inner: ConsumedResult<T>) -> LockResult<ConsumedResult<T>> {
379        self.try_write().map(|mut guard| {
380            let result = guard.deref_mut();
381            replace(result, inner)
382        })
383    }
384
385    /// similar to [parking_lot::RwLock::try_read]
386    pub fn try_read(&self) -> LockResult<RwLockReadGuard<'_, ConsumedResult<T>>> {
387        self.inner.try_read_ext()
388    }
389
390    /// similar to [parking_lot::RwLock::try_write]
391    pub fn try_write(&self) -> LockResult<RwLockWriteGuard<'_, ConsumedResult<T>>> {
392        self.inner.try_write_ext()
393    }
394
395    pub fn try_into_inner(self) -> ConsumedResult<T> {
396        self.inner.into_inner()
397    }
398
399    /// # Panics
400    ///
401    /// Panics if the internal value has already been consumed, i.e., its ownership has been moved out.
402    #[deprecated(note = "use `try_lock_inner_ref` instead")]
403    pub fn lock_inner_ref(&self) -> LockResult<MappedRwLockReadGuard<'_, T>> {
404        self.try_lock_inner_ref()
405            .map(|result| result.expect(CONSUMED_ERROR_MSG))
406    }
407
408    /// # Panics
409    ///
410    /// Panics if the internal value has already been consumed, i.e., its ownership has been moved out.
411    #[deprecated(note = "use `try_lock_inner_mut` instead")]
412    pub fn lock_inner_mut(&self) -> LockResult<MappedRwLockWriteGuard<'_, T>> {
413        self.try_lock_inner_mut()
414            .map(|result| result.expect(CONSUMED_ERROR_MSG))
415    }
416
417    /// # Panics
418    ///
419    /// Panics if the internal value has already been mutably borrowed or consumed.
420    #[deprecated(note = "use `try_lock_inner_ref` instead")]
421    pub fn inner_ref(&self) -> impl MappableDeref<'_, Target = T> {
422        self.try_lock_inner_ref()
423            .expect(LOCK_ERROR_MSG)
424            .expect(CONSUMED_ERROR_MSG)
425    }
426
427    /// # Panics
428    ///
429    /// Panics if the internal value has already been mutably borrowed or consumed.
430    #[deprecated(note = "use `try_lock_inner_mut` instead")]
431    pub fn inner_mut(&self) -> impl MappableDerefMut<'_, Target = T> {
432        self.try_lock_inner_mut()
433            .expect(LOCK_ERROR_MSG)
434            .expect(CONSUMED_ERROR_MSG)
435    }
436
437    /// # Panics
438    ///
439    /// Panics if the internal value has already been consumed, i.e., its ownership has been moved out.
440    #[deprecated(note = "use `try_into_inner` instead")]
441    pub fn into_inner(self) -> T {
442        self.try_into_inner().expect(CONSUMED_ERROR_MSG)
443    }
444}
445
446/// This trait allows you to handle [PyWrapperT0] and [PyWrapperT1] with the API of [PyWrapper]<[PyWrapperT2]>,
447/// so you can write future-compatible code.
448///
449/// # NOTE
450/// <div class="warning">
451///
452/// You must drop the returned [MappableDeref] and [MappableDerefMut], because for the implementations
453/// of [PyWrapperT1] and [PyWrapperT2], they will hold internal locks.
454///
455/// </div>
456pub trait PyWrapperSemverExt: sealed::SealedPyWrapper {
457    type Wrapped;
458
459    /// For implementations of [PyWrapper]::<[PyWrapperT1]> and ::<[PyWrapperT2]>, locks will be acquired
460    fn inner_ref_semver(
461        &self,
462    ) -> LockResult<ConsumedResult<impl MappableDeref<'_, Target = Self::Wrapped>>>;
463    /// For implementations of [PyWrapper]::<[PyWrapperT1]> and ::<[PyWrapperT2]>, locks will be acquired
464    fn inner_mut_semver(
465        &mut self,
466    ) -> LockResult<ConsumedResult<impl MappableDerefMut<'_, Target = Self::Wrapped>>>;
467    fn into_inner_semver(self) -> ConsumedResult<Self::Wrapped>;
468}
469
470impl<T> PyWrapperSemverExt for PyWrapper<PyWrapperT0<T>> {
471    type Wrapped = T;
472
473    fn inner_ref_semver(
474        &self,
475    ) -> LockResult<ConsumedResult<impl MappableDeref<'_, Target = Self::Wrapped>>> {
476        Ok(Ok(self.inner_ref()))
477    }
478
479    fn inner_mut_semver(
480        &mut self,
481    ) -> LockResult<ConsumedResult<impl MappableDerefMut<'_, Target = Self::Wrapped>>> {
482        Ok(Ok(self.inner_mut()))
483    }
484
485    fn into_inner_semver(self) -> ConsumedResult<Self::Wrapped> {
486        Ok(self.into_inner())
487    }
488}
489
490impl<T> PyWrapperSemverExt for PyWrapper<PyWrapperT1<T>> {
491    type Wrapped = T;
492
493    fn inner_ref_semver(
494        &self,
495    ) -> LockResult<ConsumedResult<impl MappableDeref<'_, Target = Self::Wrapped>>> {
496        self.lock_inner_ref().map(Ok)
497    }
498
499    fn inner_mut_semver(
500        &mut self,
501    ) -> LockResult<ConsumedResult<impl MappableDerefMut<'_, Target = Self::Wrapped>>> {
502        self.lock_inner_mut().map(Ok)
503    }
504
505    fn into_inner_semver(self) -> ConsumedResult<Self::Wrapped> {
506        Ok(self.into_inner())
507    }
508}
509
510impl<T> PyWrapperSemverExt for PyWrapper<PyWrapperT2<T>> {
511    type Wrapped = T;
512
513    fn inner_ref_semver(
514        &self,
515    ) -> LockResult<ConsumedResult<impl MappableDeref<'_, Target = Self::Wrapped>>> {
516        self.try_lock_inner_ref()
517    }
518
519    fn inner_mut_semver(
520        &mut self,
521    ) -> LockResult<ConsumedResult<impl MappableDerefMut<'_, Target = Self::Wrapped>>> {
522        self.try_lock_inner_mut()
523    }
524
525    fn into_inner_semver(self) -> ConsumedResult<Self::Wrapped> {
526        self.try_into_inner()
527    }
528}
pyo3_utils/py_wrapper.rs

pyo3_utils/
py_wrapper.rs
