use crate::impl_::pycell::PyClassObjectBaseLayout as _;
use crate::impl_::pyclass::PyClassImpl;
#[cfg(feature = "experimental-inspect")]
use crate::inspect::PyStaticExpr;
use crate::pycell::impl_::PyClassObjectLayout as _;
use crate::pycell::PyBorrowMutError;
use crate::pycell::{impl_::PyClassBorrowChecker, PyBorrowError};
use crate::pyclass::boolean_struct::False;
use crate::{ffi, Borrowed, CastError, FromPyObject, IntoPyObject, Py, PyClass, PyErr};
use std::convert::Infallible;
use std::fmt;
use std::marker::PhantomData;
use std::ops::{Deref, DerefMut};
use std::ptr::NonNull;
#[repr(transparent)]
pub struct PyClassGuard<'a, T: PyClass> {
ptr: NonNull<ffi::PyObject>,
marker: PhantomData<&'a Py<T>>,
}
impl<'a, T: PyClass> PyClassGuard<'a, T> {
pub(crate) fn try_borrow(obj: &'a Py<T>) -> Result<Self, PyBorrowError> {
Self::try_from_class_object(obj.get_class_object())
}
pub(crate) fn try_borrow_from_borrowed(
obj: Borrowed<'a, '_, T>,
) -> Result<Self, PyBorrowError> {
Self::try_from_class_object(obj.get_class_object())
}
fn try_from_class_object(obj: &'a <T as PyClassImpl>::Layout) -> Result<Self, PyBorrowError> {
obj.ensure_threadsafe();
obj.borrow_checker().try_borrow().map(|_| Self {
ptr: NonNull::from(obj).cast(),
marker: PhantomData,
})
}
pub(crate) fn as_class_object(&self) -> &'a <T as PyClassImpl>::Layout {
unsafe { self.ptr.cast().as_ref() }
}
pub fn map<F, U: ?Sized>(self, f: F) -> PyClassGuardMap<'a, U, false>
where
F: FnOnce(&T) -> &U,
{
let slf = std::mem::ManuallyDrop::new(self); PyClassGuardMap {
ptr: NonNull::from(f(&slf)),
checker: slf.as_class_object().borrow_checker(),
}
}
}
impl<'a, T> PyClassGuard<'a, T>
where
T: PyClass,
T::BaseType: PyClass,
{
pub fn as_super(&self) -> &PyClassGuard<'a, T::BaseType> {
unsafe { NonNull::from(self).cast().as_ref() }
}
pub fn into_super(self) -> PyClassGuard<'a, T::BaseType> {
let t_not_frozen = !<T::Frozen as crate::pyclass::boolean_struct::private::Boolean>::VALUE;
let u_frozen =
<<T::BaseType as PyClass>::Frozen as crate::pyclass::boolean_struct::private::Boolean>::VALUE;
if t_not_frozen && u_frozen {
self.as_super()
.as_class_object()
.borrow_checker()
.try_borrow()
.expect("this object is already borrowed");
self.as_class_object().borrow_checker().release_borrow()
};
PyClassGuard {
ptr: std::mem::ManuallyDrop::new(self).ptr,
marker: PhantomData,
}
}
}
impl<T: PyClass> Deref for PyClassGuard<'_, T> {
type Target = T;
#[inline]
fn deref(&self) -> &T {
unsafe { &*self.as_class_object().get_ptr().cast_const() }
}
}
impl<'a, 'py, T: PyClass> FromPyObject<'a, 'py> for PyClassGuard<'a, T> {
type Error = PyClassGuardError<'a, 'py>;
#[cfg(feature = "experimental-inspect")]
const INPUT_TYPE: PyStaticExpr = T::TYPE_HINT;
fn extract(obj: Borrowed<'a, 'py, crate::PyAny>) -> Result<Self, Self::Error> {
Self::try_from_class_object(
obj.cast::<T>()
.map_err(|e| PyClassGuardError(Some(e)))?
.get_class_object(),
)
.map_err(|_| PyClassGuardError(None))
}
}
impl<'a, 'py, T: PyClass> IntoPyObject<'py> for PyClassGuard<'a, T> {
type Target = T;
type Output = Borrowed<'a, 'py, T>;
type Error = Infallible;
#[cfg(feature = "experimental-inspect")]
const OUTPUT_TYPE: PyStaticExpr = T::TYPE_HINT;
#[inline]
fn into_pyobject(self, py: crate::Python<'py>) -> Result<Self::Output, Self::Error> {
(&self).into_pyobject(py)
}
}
impl<'a, 'py, T: PyClass> IntoPyObject<'py> for &PyClassGuard<'a, T> {
type Target = T;
type Output = Borrowed<'a, 'py, T>;
type Error = Infallible;
#[cfg(feature = "experimental-inspect")]
const OUTPUT_TYPE: PyStaticExpr = T::TYPE_HINT;
#[inline]
fn into_pyobject(self, py: crate::Python<'py>) -> Result<Self::Output, Self::Error> {
unsafe { Ok(Borrowed::from_non_null(py, self.ptr).cast_unchecked()) }
}
}
impl<T: PyClass> Drop for PyClassGuard<'_, T> {
fn drop(&mut self) {
self.as_class_object().borrow_checker().release_borrow()
}
}
#[cfg(feature = "nightly")]
unsafe impl<T: PyClass> crate::marker::Ungil for PyClassGuard<'_, T> {}
unsafe impl<T: PyClass + Sync> Send for PyClassGuard<'_, T> {}
unsafe impl<T: PyClass + Sync> Sync for PyClassGuard<'_, T> {}
pub struct PyClassGuardError<'a, 'py>(pub(crate) Option<CastError<'a, 'py>>);
impl fmt::Debug for PyClassGuardError<'_, '_> {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
if let Some(e) = &self.0 {
write!(f, "{e:?}")
} else {
write!(f, "{:?}", PyBorrowError::new())
}
}
}
impl fmt::Display for PyClassGuardError<'_, '_> {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
if let Some(e) = &self.0 {
write!(f, "{e}")
} else {
write!(f, "{}", PyBorrowError::new())
}
}
}
impl From<PyClassGuardError<'_, '_>> for PyErr {
fn from(value: PyClassGuardError<'_, '_>) -> Self {
if let Some(e) = value.0 {
e.into()
} else {
PyBorrowError::new().into()
}
}
}
#[repr(transparent)]
pub struct PyClassGuardMut<'a, T: PyClass<Frozen = False>> {
ptr: NonNull<ffi::PyObject>,
marker: PhantomData<&'a Py<T>>,
}
impl<'a, T: PyClass<Frozen = False>> PyClassGuardMut<'a, T> {
pub(crate) fn try_borrow_mut(obj: &'a Py<T>) -> Result<Self, PyBorrowMutError> {
Self::try_from_class_object(obj.get_class_object())
}
pub(crate) fn try_borrow_mut_from_borrowed(
obj: Borrowed<'a, '_, T>,
) -> Result<Self, PyBorrowMutError> {
Self::try_from_class_object(obj.get_class_object())
}
fn try_from_class_object(
obj: &'a <T as PyClassImpl>::Layout,
) -> Result<Self, PyBorrowMutError> {
obj.ensure_threadsafe();
obj.borrow_checker().try_borrow_mut().map(|_| Self {
ptr: NonNull::from(obj).cast(),
marker: PhantomData,
})
}
pub(crate) fn as_class_object(&self) -> &'a <T as PyClassImpl>::Layout {
unsafe { self.ptr.cast().as_ref() }
}
pub fn map<F, U: ?Sized>(self, f: F) -> PyClassGuardMap<'a, U, true>
where
F: FnOnce(&mut T) -> &mut U,
{
let mut slf = std::mem::ManuallyDrop::new(self); PyClassGuardMap {
ptr: NonNull::from(f(&mut slf)),
checker: slf.as_class_object().borrow_checker(),
}
}
}
impl<'a, T> PyClassGuardMut<'a, T>
where
T: PyClass<Frozen = False>,
T::BaseType: PyClass<Frozen = False>,
{
pub fn as_super(&mut self) -> &mut PyClassGuardMut<'a, T::BaseType> {
unsafe { NonNull::from(self).cast().as_mut() }
}
pub fn into_super(self) -> PyClassGuardMut<'a, T::BaseType> {
PyClassGuardMut {
ptr: std::mem::ManuallyDrop::new(self).ptr,
marker: PhantomData,
}
}
}
impl<T: PyClass<Frozen = False>> Deref for PyClassGuardMut<'_, T> {
type Target = T;
#[inline]
fn deref(&self) -> &T {
unsafe { &*self.as_class_object().get_ptr().cast_const() }
}
}
impl<T: PyClass<Frozen = False>> DerefMut for PyClassGuardMut<'_, T> {
#[inline]
fn deref_mut(&mut self) -> &mut T {
unsafe { &mut *self.as_class_object().get_ptr() }
}
}
impl<'a, 'py, T: PyClass<Frozen = False>> FromPyObject<'a, 'py> for PyClassGuardMut<'a, T> {
type Error = PyClassGuardMutError<'a, 'py>;
#[cfg(feature = "experimental-inspect")]
const INPUT_TYPE: PyStaticExpr = T::TYPE_HINT;
fn extract(obj: Borrowed<'a, 'py, crate::PyAny>) -> Result<Self, Self::Error> {
Self::try_from_class_object(
obj.cast::<T>()
.map_err(|e| PyClassGuardMutError(Some(e)))?
.get_class_object(),
)
.map_err(|_| PyClassGuardMutError(None))
}
}
impl<'a, 'py, T: PyClass<Frozen = False>> IntoPyObject<'py> for PyClassGuardMut<'a, T> {
type Target = T;
type Output = Borrowed<'a, 'py, T>;
type Error = Infallible;
#[cfg(feature = "experimental-inspect")]
const OUTPUT_TYPE: PyStaticExpr = T::TYPE_HINT;
#[inline]
fn into_pyobject(self, py: crate::Python<'py>) -> Result<Self::Output, Self::Error> {
(&self).into_pyobject(py)
}
}
impl<'a, 'py, T: PyClass<Frozen = False>> IntoPyObject<'py> for &PyClassGuardMut<'a, T> {
type Target = T;
type Output = Borrowed<'a, 'py, T>;
type Error = Infallible;
#[cfg(feature = "experimental-inspect")]
const OUTPUT_TYPE: PyStaticExpr = T::TYPE_HINT;
#[inline]
fn into_pyobject(self, py: crate::Python<'py>) -> Result<Self::Output, Self::Error> {
unsafe { Ok(Borrowed::from_non_null(py, self.ptr).cast_unchecked()) }
}
}
impl<T: PyClass<Frozen = False>> Drop for PyClassGuardMut<'_, T> {
fn drop(&mut self) {
self.as_class_object().borrow_checker().release_borrow_mut()
}
}
#[cfg(feature = "nightly")]
unsafe impl<T: PyClass<Frozen = False>> crate::marker::Ungil for PyClassGuardMut<'_, T> {}
unsafe impl<T: PyClass<Frozen = False> + Send + Sync> Send for PyClassGuardMut<'_, T> {}
unsafe impl<T: PyClass<Frozen = False> + Sync> Sync for PyClassGuardMut<'_, T> {}
pub struct PyClassGuardMutError<'a, 'py>(pub(crate) Option<CastError<'a, 'py>>);
impl fmt::Debug for PyClassGuardMutError<'_, '_> {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
if let Some(e) = &self.0 {
write!(f, "{e:?}")
} else {
write!(f, "{:?}", PyBorrowMutError::new())
}
}
}
impl fmt::Display for PyClassGuardMutError<'_, '_> {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
if let Some(e) = &self.0 {
write!(f, "{e}")
} else {
write!(f, "{}", PyBorrowMutError::new())
}
}
}
impl From<PyClassGuardMutError<'_, '_>> for PyErr {
fn from(value: PyClassGuardMutError<'_, '_>) -> Self {
if let Some(e) = value.0 {
e.into()
} else {
PyBorrowMutError::new().into()
}
}
}
pub struct PyClassGuardMap<'a, U: ?Sized, const MUT: bool> {
ptr: NonNull<U>,
checker: &'a dyn PyClassBorrowChecker,
}
impl<U: ?Sized, const MUT: bool> Deref for PyClassGuardMap<'_, U, MUT> {
type Target = U;
fn deref(&self) -> &U {
unsafe { self.ptr.as_ref() }
}
}
impl<U: ?Sized> DerefMut for PyClassGuardMap<'_, U, true> {
fn deref_mut(&mut self) -> &mut Self::Target {
unsafe { self.ptr.as_mut() }
}
}
impl<U: ?Sized, const MUT: bool> Drop for PyClassGuardMap<'_, U, MUT> {
fn drop(&mut self) {
if MUT {
self.checker.release_borrow_mut();
} else {
self.checker.release_borrow();
}
}
}
#[cfg(test)]
#[cfg(feature = "macros")]
mod tests {
use super::{PyClassGuard, PyClassGuardMut};
use crate::{types::PyAnyMethods as _, Bound, IntoPyObject as _, Py, PyErr, Python};
#[test]
fn test_into_frozen_super_released_borrow() {
#[crate::pyclass]
#[pyo3(crate = "crate", subclass, frozen)]
struct BaseClass {}
#[crate::pyclass]
#[pyo3(crate = "crate", extends=BaseClass, subclass)]
struct SubClass {}
#[crate::pymethods]
#[pyo3(crate = "crate")]
impl SubClass {
#[new]
fn new(py: Python<'_>) -> Py<SubClass> {
let init = crate::PyClassInitializer::from(BaseClass {}).add_subclass(SubClass {});
Py::new(py, init).expect("allocation error")
}
}
Python::attach(|py| {
let obj = SubClass::new(py);
drop(PyClassGuard::try_borrow(&obj).unwrap().into_super());
assert!(PyClassGuardMut::try_borrow_mut(&obj).is_ok());
})
}
#[test]
fn test_into_frozen_super_mutable_base_holds_borrow() {
#[crate::pyclass]
#[pyo3(crate = "crate", subclass)]
struct BaseClass {}
#[crate::pyclass]
#[pyo3(crate = "crate", extends=BaseClass, subclass, frozen)]
struct SubClass {}
#[crate::pyclass]
#[pyo3(crate = "crate", extends=SubClass, subclass)]
struct SubSubClass {}
#[crate::pymethods]
#[pyo3(crate = "crate")]
impl SubSubClass {
#[new]
fn new(py: Python<'_>) -> Py<SubSubClass> {
let init = crate::PyClassInitializer::from(BaseClass {})
.add_subclass(SubClass {})
.add_subclass(SubSubClass {});
Py::new(py, init).expect("allocation error")
}
}
Python::attach(|py| {
let obj = SubSubClass::new(py);
let _super_borrow = PyClassGuard::try_borrow(&obj).unwrap().into_super();
assert!(PyClassGuardMut::try_borrow_mut(&obj).is_err());
})
}
#[crate::pyclass]
#[pyo3(crate = "crate", subclass)]
struct BaseClass {
val1: usize,
}
#[crate::pyclass]
#[pyo3(crate = "crate", extends=BaseClass, subclass)]
struct SubClass {
val2: usize,
}
#[crate::pyclass]
#[pyo3(crate = "crate", extends=SubClass)]
struct SubSubClass {
#[pyo3(get)]
val3: usize,
}
#[crate::pymethods]
#[pyo3(crate = "crate")]
impl SubSubClass {
#[new]
fn new(py: Python<'_>) -> Py<SubSubClass> {
let init = crate::PyClassInitializer::from(BaseClass { val1: 10 })
.add_subclass(SubClass { val2: 15 })
.add_subclass(SubSubClass { val3: 20 });
Py::new(py, init).expect("allocation error")
}
fn get_values(self_: PyClassGuard<'_, Self>) -> (usize, usize, usize) {
let val1 = self_.as_super().as_super().val1;
let val2 = self_.as_super().val2;
(val1, val2, self_.val3)
}
fn double_values(mut self_: PyClassGuardMut<'_, Self>) {
self_.as_super().as_super().val1 *= 2;
self_.as_super().val2 *= 2;
self_.val3 *= 2;
}
fn __add__<'a>(
mut slf: PyClassGuardMut<'a, Self>,
other: PyClassGuard<'a, Self>,
) -> PyClassGuardMut<'a, Self> {
slf.val3 += other.val3;
slf
}
fn __rsub__<'a>(
slf: PyClassGuard<'a, Self>,
mut other: PyClassGuardMut<'a, Self>,
) -> PyClassGuardMut<'a, Self> {
other.val3 -= slf.val3;
other
}
}
#[test]
fn test_pyclassguard_into_pyobject() {
Python::attach(|py| {
let class = Py::new(py, BaseClass { val1: 42 })?;
let guard = PyClassGuard::try_borrow(&class).unwrap();
let new_ref = (&guard).into_pyobject(py)?;
assert!(new_ref.is(&class));
let new = guard.into_pyobject(py)?;
assert!(new.is(&class));
Ok::<_, PyErr>(())
})
.unwrap();
}
#[test]
fn test_pyclassguardmut_into_pyobject() {
Python::attach(|py| {
let class = Py::new(py, BaseClass { val1: 42 })?;
let guard = PyClassGuardMut::try_borrow_mut(&class).unwrap();
let new_ref = (&guard).into_pyobject(py)?;
assert!(new_ref.is(&class));
let new = guard.into_pyobject(py)?;
assert!(new.is(&class));
Ok::<_, PyErr>(())
})
.unwrap();
}
#[test]
fn test_pyclassguard_as_super() {
Python::attach(|py| {
let obj = SubSubClass::new(py).into_bound(py);
let pyref = PyClassGuard::try_borrow(obj.as_unbound()).unwrap();
assert_eq!(pyref.as_super().as_super().val1, 10);
assert_eq!(pyref.as_super().val2, 15);
assert_eq!(pyref.val3, 20);
assert_eq!(SubSubClass::get_values(pyref), (10, 15, 20));
});
}
#[test]
fn test_pyclassguardmut_as_super() {
Python::attach(|py| {
let obj = SubSubClass::new(py).into_bound(py);
assert_eq!(
SubSubClass::get_values(PyClassGuard::try_borrow(obj.as_unbound()).unwrap()),
(10, 15, 20)
);
{
let mut pyrefmut = PyClassGuardMut::try_borrow_mut(obj.as_unbound()).unwrap();
assert_eq!(pyrefmut.as_super().as_super().val1, 10);
pyrefmut.as_super().as_super().val1 -= 5;
pyrefmut.as_super().val2 -= 5;
pyrefmut.val3 -= 5;
}
assert_eq!(
SubSubClass::get_values(PyClassGuard::try_borrow(obj.as_unbound()).unwrap()),
(5, 10, 15)
);
SubSubClass::double_values(PyClassGuardMut::try_borrow_mut(obj.as_unbound()).unwrap());
assert_eq!(
SubSubClass::get_values(PyClassGuard::try_borrow(obj.as_unbound()).unwrap()),
(10, 20, 30)
);
});
}
#[test]
fn test_extract_guard() {
Python::attach(|py| {
let obj1 = SubSubClass::new(py);
let obj2 = SubSubClass::new(py);
crate::py_run!(py, obj1 obj2, "assert ((obj1 + obj2) - obj2).val3 == obj1.val3");
});
}
#[test]
fn test_pyclassguards_in_python() {
Python::attach(|py| {
let obj = SubSubClass::new(py);
crate::py_run!(py, obj, "assert obj.get_values() == (10, 15, 20)");
crate::py_run!(py, obj, "assert obj.double_values() is None");
crate::py_run!(py, obj, "assert obj.get_values() == (20, 30, 40)");
});
}
#[crate::pyclass]
#[pyo3(crate = "crate")]
pub struct MyClass {
data: [i32; 100],
}
#[test]
fn test_pyclassguard_map() {
Python::attach(|py| {
let obj = Bound::new(py, MyClass { data: [0; 100] })?;
let data = PyClassGuard::try_borrow(obj.as_unbound())?.map(|c| &c.data);
assert_eq!(data[0], 0);
assert!(obj.try_borrow_mut().is_err()); drop(data);
assert!(obj.try_borrow_mut().is_ok()); Ok::<_, PyErr>(())
})
.unwrap()
}
#[test]
fn test_pyclassguardmut_map() {
Python::attach(|py| {
let obj = Bound::new(py, MyClass { data: [0; 100] })?;
let mut data =
PyClassGuardMut::try_borrow_mut(obj.as_unbound())?.map(|c| c.data.as_mut_slice());
assert_eq!(data[0], 0);
data[0] = 5;
assert_eq!(data[0], 5);
assert!(obj.try_borrow_mut().is_err()); drop(data);
assert!(obj.try_borrow_mut().is_ok()); Ok::<_, PyErr>(())
})
.unwrap()
}
#[test]
fn test_pyclassguard_map_unrelated() {
use crate::types::{PyString, PyStringMethods};
Python::attach(|py| {
let obj = Bound::new(py, MyClass { data: [0; 100] })?;
let string = PyString::new(py, "pyo3");
let refmap = PyClassGuard::try_borrow(obj.as_unbound())?.map(|_| &string);
assert_eq!(refmap.to_cow()?, "pyo3");
Ok::<_, PyErr>(())
})
.unwrap()
}
}