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        match self.inner.as_ref() {
282            Ok(inner) => inner,
283            // this arm cannot be omitted because of the reference,
284            // see: <https://blog.rust-lang.org/2024/10/17/Rust-1.82.0.html#omitting-empty-types-in-pattern-matching>
285            Err(infallible) => match *infallible {},
286        }
287    }
288
289    #[inline]
290    pub fn inner_mut(&mut self) -> impl MappableDerefMut<'_, Target = T> {
291        // TODO, FIXME: use [Result::into_ok] instead (unstable for now)
292        match self.inner.as_mut() {
293            Ok(inner) => inner,
294            Err(infallible) => match *infallible {},
295        }
296    }
297
298    #[inline]
299    pub fn into_inner(self) -> T {
300        // TODO, FIXME: use [Result::into_ok] instead (unstable for now)
301        let Ok(inner) = self.inner;
302        inner
303    }
304}
305
306impl<T> PyWrapper<PyWrapperT1<T>> {
307    #[inline]
308    pub fn new1(inner: T) -> Self {
309        Self {
310            inner: RwLock::new(Ok(inner)),
311        }
312    }
313
314    pub fn lock_inner_ref(&self) -> LockResult<MappedRwLockReadGuard<'_, T>> {
315        self.inner
316            .try_read_ext()
317            // TODO, FIXME: use [Result::into_ok] instead (unstable for now)
318            .map(|guard| {
319                RwLockReadGuard::map(guard, |inner| match inner.as_ref() {
320                    Ok(inner) => inner,
321                    Err(infallible) => match *infallible {},
322                })
323            })
324    }
325
326    pub fn lock_inner_mut(&self) -> LockResult<MappedRwLockWriteGuard<'_, T>> {
327        self.inner
328            .try_write_ext()
329            // TODO, FIXME: use [Result::into_ok] instead (unstable for now)
330            .map(|guard| {
331                RwLockWriteGuard::map(guard, |inner| match inner.as_mut() {
332                    Ok(inner) => inner,
333                    Err(infallible) => match *infallible {},
334                })
335            })
336    }
337
338    pub fn into_inner(self) -> T {
339        // TODO, FIXME: use [Result::into_ok] instead (unstable for now)
340        let Ok(inner) = self.inner.into_inner();
341        inner
342    }
343
344    /// # Panics
345    ///
346    /// Panics if the internal value has already been mutably borrowed.
347    #[deprecated(note = "use `lock_inner_ref` instead")]
348    pub fn inner_ref(&self) -> impl MappableDeref<'_, Target = T> {
349        self.lock_inner_ref().expect(LOCK_ERROR_MSG)
350    }
351
352    /// # Panics
353    ///
354    /// Panics if the internal value has already been mutably borrowed.
355    #[deprecated(note = "use `lock_inner_mut` instead")]
356    pub fn inner_mut(&self) -> impl MappableDerefMut<'_, Target = T> {
357        self.lock_inner_mut().expect(LOCK_ERROR_MSG)
358    }
359}
360
361impl<T> PyWrapper<PyWrapperT2<T>> {
362    #[inline]
363    pub fn new2(inner: T) -> Self {
364        Self {
365            inner: RwLock::new(Ok(inner)),
366        }
367    }
368
369    pub fn try_lock_inner_ref(&self) -> LockResult<ConsumedResult<MappedRwLockReadGuard<'_, T>>> {
370        self.try_read().map(|guard| {
371            if guard.is_err() {
372                Err(ConsumedError)
373            } else {
374                // PEFR: it's ok to use [unwrap_unchecked], but i dont like unsafe block
375                Ok(RwLockReadGuard::map(guard, |inner| inner.as_ref().unwrap()))
376            }
377        })
378    }
379
380    pub fn try_lock_inner_mut(&self) -> LockResult<ConsumedResult<MappedRwLockWriteGuard<'_, T>>> {
381        self.try_write().map(|guard| {
382            if guard.is_err() {
383                Err(ConsumedError)
384            } else {
385                // PEFR: it's ok to use [unwrap_unchecked], but i dont like unsafe block
386                Ok(RwLockWriteGuard::map(guard, |inner| {
387                    inner.as_mut().unwrap()
388                }))
389            }
390        })
391    }
392
393    pub fn try_take_inner(&self) -> LockResult<ConsumedResult<T>> {
394        self.try_replace_inner(Err(ConsumedError))
395    }
396
397    /// similar to [std::mem::replace]
398    pub fn try_replace_inner(&self, inner: ConsumedResult<T>) -> LockResult<ConsumedResult<T>> {
399        self.try_write().map(|mut guard| {
400            let result = guard.deref_mut();
401            replace(result, inner)
402        })
403    }
404
405    /// similar to [parking_lot::RwLock::try_read]
406    pub fn try_read(&self) -> LockResult<RwLockReadGuard<'_, ConsumedResult<T>>> {
407        self.inner.try_read_ext()
408    }
409
410    /// similar to [parking_lot::RwLock::try_write]
411    pub fn try_write(&self) -> LockResult<RwLockWriteGuard<'_, ConsumedResult<T>>> {
412        self.inner.try_write_ext()
413    }
414
415    pub fn try_into_inner(self) -> ConsumedResult<T> {
416        self.inner.into_inner()
417    }
418
419    /// # Panics
420    ///
421    /// Panics if the internal value has already been consumed, i.e., its ownership has been moved out.
422    #[deprecated(note = "use `try_lock_inner_ref` instead")]
423    pub fn lock_inner_ref(&self) -> LockResult<MappedRwLockReadGuard<'_, T>> {
424        self.try_lock_inner_ref()
425            .map(|result| result.expect(CONSUMED_ERROR_MSG))
426    }
427
428    /// # Panics
429    ///
430    /// Panics if the internal value has already been consumed, i.e., its ownership has been moved out.
431    #[deprecated(note = "use `try_lock_inner_mut` instead")]
432    pub fn lock_inner_mut(&self) -> LockResult<MappedRwLockWriteGuard<'_, T>> {
433        self.try_lock_inner_mut()
434            .map(|result| result.expect(CONSUMED_ERROR_MSG))
435    }
436
437    /// # Panics
438    ///
439    /// Panics if the internal value has already been mutably borrowed or consumed.
440    #[deprecated(note = "use `try_lock_inner_ref` instead")]
441    pub fn inner_ref(&self) -> impl MappableDeref<'_, Target = T> {
442        self.try_lock_inner_ref()
443            .expect(LOCK_ERROR_MSG)
444            .expect(CONSUMED_ERROR_MSG)
445    }
446
447    /// # Panics
448    ///
449    /// Panics if the internal value has already been mutably borrowed or consumed.
450    #[deprecated(note = "use `try_lock_inner_mut` instead")]
451    pub fn inner_mut(&self) -> impl MappableDerefMut<'_, Target = T> {
452        self.try_lock_inner_mut()
453            .expect(LOCK_ERROR_MSG)
454            .expect(CONSUMED_ERROR_MSG)
455    }
456
457    /// # Panics
458    ///
459    /// Panics if the internal value has already been consumed, i.e., its ownership has been moved out.
460    #[deprecated(note = "use `try_into_inner` instead")]
461    pub fn into_inner(self) -> T {
462        self.try_into_inner().expect(CONSUMED_ERROR_MSG)
463    }
464}
465
466/// This trait allows you to handle [PyWrapperT0] and [PyWrapperT1] with the API of [PyWrapper]<[PyWrapperT2]>,
467/// so you can write future-compatible code.
468///
469/// # NOTE
470/// <div class="warning">
471///
472/// You must drop the returned [MappableDeref] and [MappableDerefMut], because for the implementations
473/// of [PyWrapperT1] and [PyWrapperT2], they will hold internal locks.
474///
475/// </div>
476pub trait PyWrapperSemverExt: sealed::SealedPyWrapper {
477    type Wrapped;
478
479    /// For implementations of [PyWrapper]::<[PyWrapperT1]> and ::<[PyWrapperT2]>, locks will be acquired
480    fn inner_ref_semver(
481        &self,
482    ) -> LockResult<ConsumedResult<impl MappableDeref<'_, Target = Self::Wrapped>>>;
483    /// For implementations of [PyWrapper]::<[PyWrapperT1]> and ::<[PyWrapperT2]>, locks will be acquired
484    fn inner_mut_semver(
485        &mut self,
486    ) -> LockResult<ConsumedResult<impl MappableDerefMut<'_, Target = Self::Wrapped>>>;
487    fn into_inner_semver(self) -> ConsumedResult<Self::Wrapped>;
488}
489
490impl<T> PyWrapperSemverExt for PyWrapper<PyWrapperT0<T>> {
491    type Wrapped = T;
492
493    fn inner_ref_semver(
494        &self,
495    ) -> LockResult<ConsumedResult<impl MappableDeref<'_, Target = Self::Wrapped>>> {
496        Ok(Ok(self.inner_ref()))
497    }
498
499    fn inner_mut_semver(
500        &mut self,
501    ) -> LockResult<ConsumedResult<impl MappableDerefMut<'_, Target = Self::Wrapped>>> {
502        Ok(Ok(self.inner_mut()))
503    }
504
505    fn into_inner_semver(self) -> ConsumedResult<Self::Wrapped> {
506        Ok(self.into_inner())
507    }
508}
509
510impl<T> PyWrapperSemverExt for PyWrapper<PyWrapperT1<T>> {
511    type Wrapped = T;
512
513    fn inner_ref_semver(
514        &self,
515    ) -> LockResult<ConsumedResult<impl MappableDeref<'_, Target = Self::Wrapped>>> {
516        self.lock_inner_ref().map(Ok)
517    }
518
519    fn inner_mut_semver(
520        &mut self,
521    ) -> LockResult<ConsumedResult<impl MappableDerefMut<'_, Target = Self::Wrapped>>> {
522        self.lock_inner_mut().map(Ok)
523    }
524
525    fn into_inner_semver(self) -> ConsumedResult<Self::Wrapped> {
526        Ok(self.into_inner())
527    }
528}
529
530impl<T> PyWrapperSemverExt for PyWrapper<PyWrapperT2<T>> {
531    type Wrapped = T;
532
533    fn inner_ref_semver(
534        &self,
535    ) -> LockResult<ConsumedResult<impl MappableDeref<'_, Target = Self::Wrapped>>> {
536        self.try_lock_inner_ref()
537    }
538
539    fn inner_mut_semver(
540        &mut self,
541    ) -> LockResult<ConsumedResult<impl MappableDerefMut<'_, Target = Self::Wrapped>>> {
542        self.try_lock_inner_mut()
543    }
544
545    fn into_inner_semver(self) -> ConsumedResult<Self::Wrapped> {
546        self.try_into_inner()
547    }
548}
pyo3_utils/py_wrapper.rs

pyo3_utils/
py_wrapper.rs
