From 35f7f1a78c5914567bdd63c179bf8884691fe83a Mon Sep 17 00:00:00 2001 From: David Hewitt Date: Thu, 14 Dec 2023 19:07:47 +0000 Subject: [PATCH] use Py2Borrowed to make PyBytesMethods slightly nicer --- src/ffi_ptr_ext.rs | 45 +++++++++++++++++++- src/instance.rs | 94 ++++++++++++++++++++++++++++++++++++++++++ src/types/bytearray.rs | 32 +++++++++----- src/types/bytes.rs | 19 +++++---- 4 files changed, 171 insertions(+), 19 deletions(-) diff --git a/src/ffi_ptr_ext.rs b/src/ffi_ptr_ext.rs index 34f95cd2..0034ea85 100644 --- a/src/ffi_ptr_ext.rs +++ b/src/ffi_ptr_ext.rs @@ -1,4 +1,8 @@ -use crate::{ffi, instance::Py2, PyAny, PyResult, Python}; +use crate::{ + ffi, + instance::{Py2, Py2Borrowed}, + PyAny, PyResult, Python, +}; mod sealed { use super::*; @@ -13,6 +17,24 @@ use sealed::Sealed; pub(crate) trait FfiPtrExt: Sealed { unsafe fn assume_owned_or_err(self, py: Python<'_>) -> PyResult>; unsafe fn assume_owned(self, py: Python<'_>) -> Py2<'_, PyAny>; + + /// Assumes this pointer is borrowed from a parent object. + /// + /// Warning: the lifetime `'a` is not bounded by the function arguments; the caller is + /// responsible to ensure this is tied to some appropriate lifetime. + unsafe fn assume_borrowed_or_err<'a>( + self, + py: Python<'_>, + ) -> PyResult>; + + /// Same as `assume_borrowed_or_err`, but doesn't fetch an error on NULL. + unsafe fn assume_borrowed_or_opt<'a>( + self, + py: Python<'_>, + ) -> Option>; + + /// Same as `assume_borrowed_or_err`, but panics on NULL. + unsafe fn assume_borrowed<'a>(self, py: Python<'_>) -> Py2Borrowed<'a, '_, PyAny>; } impl FfiPtrExt for *mut ffi::PyObject { @@ -25,4 +47,25 @@ impl FfiPtrExt for *mut ffi::PyObject { unsafe fn assume_owned(self, py: Python<'_>) -> Py2<'_, PyAny> { Py2::from_owned_ptr(py, self) } + + #[inline] + unsafe fn assume_borrowed_or_err<'a>( + self, + py: Python<'_>, + ) -> PyResult> { + Py2Borrowed::from_ptr_or_err(py, self) + } + + #[inline] + unsafe fn assume_borrowed_or_opt<'a>( + self, + py: Python<'_>, + ) -> Option> { + Py2Borrowed::from_ptr_or_opt(py, self) + } + + #[inline] + unsafe fn assume_borrowed<'a>(self, py: Python<'_>) -> Py2Borrowed<'a, '_, PyAny> { + Py2Borrowed::from_ptr(py, self) + } } diff --git a/src/instance.rs b/src/instance.rs index 3ec2a3f8..7568b28f 100644 --- a/src/instance.rs +++ b/src/instance.rs @@ -217,6 +217,96 @@ unsafe impl AsPyPointer for Py2<'_, T> { } } +/// A borrowed equivalent to `Py2`. +/// +/// The advantage of this over `&Py2` is that it avoids the need to have a pointer-to-pointer, as Py2 +/// is already a pointer to an `ffi::PyObject``. +#[repr(transparent)] +pub(crate) struct Py2Borrowed<'a, 'py, T>( + NonNull, + PhantomData<&'a Py>, + Python<'py>, +); + +impl<'a, 'py> Py2Borrowed<'a, 'py, PyAny> { + /// # Safety + /// This is similar to `std::slice::from_raw_parts`, the lifetime `'a` is completely defined by + /// the caller and it's the caller's responsibility to ensure that the reference this is + /// derived from is valid for the lifetime `'a`. + pub(crate) unsafe fn from_ptr_or_err( + py: Python<'py>, + ptr: *mut ffi::PyObject, + ) -> PyResult { + NonNull::new(ptr).map_or_else( + || Err(PyErr::fetch(py)), + |ptr| Ok(Self(ptr, PhantomData, py)), + ) + } + + /// # Safety + /// This is similar to `std::slice::from_raw_parts`, the lifetime `'a` is completely defined by + /// the caller and it's the caller's responsibility to ensure that the reference this is + /// derived from is valid for the lifetime `'a`. + pub(crate) unsafe fn from_ptr_or_opt(py: Python<'py>, ptr: *mut ffi::PyObject) -> Option { + NonNull::new(ptr).map(|ptr| Self(ptr, PhantomData, py)) + } + + /// # Safety + /// This is similar to `std::slice::from_raw_parts`, the lifetime `'a` is completely defined by + /// the caller and it's the caller's responsibility to ensure that the reference this is + /// derived from is valid for the lifetime `'a`. + pub(crate) unsafe fn from_ptr(py: Python<'py>, ptr: *mut ffi::PyObject) -> Self { + Self( + NonNull::new(ptr).unwrap_or_else(|| crate::err::panic_after_error(py)), + PhantomData, + py, + ) + } +} + +impl<'a, 'py, T> From<&'a Py2<'py, T>> for Py2Borrowed<'a, 'py, T> { + /// Create borrow on a Py2 + fn from(instance: &'a Py2<'py, T>) -> Self { + Self( + unsafe { NonNull::new_unchecked(instance.as_ptr()) }, + PhantomData, + instance.py(), + ) + } +} + +impl<'py, T> Py2Borrowed<'py, 'py, T> +where + T: HasPyGilRef, +{ + pub(crate) fn from_gil_ref(gil_ref: &'py T::AsRefTarget) -> Self { + // Safety: &'py T::AsRefTarget is expected to be a Python pointer, + // so &'py T::AsRefTarget has the same layout as Self. + unsafe { std::mem::transmute(gil_ref) } + } + + // pub(crate) fn into_gil_ref(self) -> &'py T::AsRefTarget { + // // Safety: self is a borrow over `'py`. + // unsafe { self.py().from_borrowed_ptr(self.0.as_ptr()) } + // } +} + +impl std::fmt::Debug for Py2Borrowed<'_, '_, T> { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + Py2::fmt(self, f) + } +} + +impl<'py, T> Deref for Py2Borrowed<'_, 'py, T> { + type Target = Py2<'py, T>; + + #[inline] + fn deref(&self) -> &Py2<'py, T> { + // safety: Py2 has the same layout as NonNull + unsafe { &*(&self.0 as *const _ as *const Py2<'py, T>) } + } +} + /// A GIL-independent reference to an object allocated on the Python heap. /// /// This type does not auto-dereference to the inner object because you must prove you hold the GIL to access it. @@ -727,6 +817,10 @@ impl Py { Py2(py, ManuallyDrop::new(self)) } + pub(crate) fn attach_borrow<'a, 'py>(&'a self, py: Python<'py>) -> Py2Borrowed<'a, 'py, T> { + Py2Borrowed(self.0, PhantomData, py) + } + /// Returns whether `self` and `other` point to the same object. To compare /// the equality of two objects (the `==` operator), use [`eq`](PyAny::eq). /// diff --git a/src/types/bytearray.rs b/src/types/bytearray.rs index 885490cb..e7f02179 100644 --- a/src/types/bytearray.rs +++ b/src/types/bytearray.rs @@ -1,5 +1,5 @@ use crate::err::{PyErr, PyResult}; -use crate::instance::Py2; +use crate::instance::{Py2, Py2Borrowed}; use crate::{ffi, AsPyPointer, Py, PyAny, Python}; use std::os::raw::c_char; use std::slice; @@ -188,9 +188,7 @@ impl PyByteArray { /// } /// ``` pub unsafe fn as_bytes(&self) -> &[u8] { - let slice = Py2::borrowed_from_gil_ref(&self).as_bytes(); - // SAFETY: &self guarantees the reference is alive long enough - unsafe { slice::from_raw_parts(slice.as_ptr(), slice.len()) } + Py2Borrowed::from_gil_ref(self).as_bytes() } /// Extracts a mutable slice of the `ByteArray`'s entire buffer. @@ -202,9 +200,7 @@ impl PyByteArray { /// apply to this function as well. #[allow(clippy::mut_from_ref)] pub unsafe fn as_bytes_mut(&self) -> &mut [u8] { - let slice = Py2::borrowed_from_gil_ref(&self).as_bytes_mut(); - // SAFETY: &self guarantees the reference is alive long enough - unsafe { slice::from_raw_parts_mut(slice.as_mut_ptr(), slice.len()) } + Py2Borrowed::from_gil_ref(self).as_bytes_mut() } /// Copies the contents of the bytearray to a Rust vector. @@ -404,16 +400,16 @@ impl<'py> PyByteArrayMethods<'py> for Py2<'py, PyByteArray> { } fn data(&self) -> *mut u8 { - unsafe { ffi::PyByteArray_AsString(self.as_ptr()).cast() } + Py2Borrowed::from(self).data() } unsafe fn as_bytes(&self) -> &[u8] { - slice::from_raw_parts(self.data(), self.len()) + Py2Borrowed::from(self).as_bytes() } #[allow(clippy::mut_from_ref)] unsafe fn as_bytes_mut(&self) -> &mut [u8] { - slice::from_raw_parts_mut(self.data(), self.len()) + Py2Borrowed::from(self).as_bytes_mut() } fn to_vec(&self) -> Vec { @@ -432,6 +428,22 @@ impl<'py> PyByteArrayMethods<'py> for Py2<'py, PyByteArray> { } } +impl<'a> Py2Borrowed<'a, '_, PyByteArray> { + fn data(&self) -> *mut u8 { + unsafe { ffi::PyByteArray_AsString(self.as_ptr()).cast() } + } + + #[allow(clippy::wrong_self_convention)] + unsafe fn as_bytes(self) -> &'a [u8] { + slice::from_raw_parts(self.data(), self.len()) + } + + #[allow(clippy::wrong_self_convention)] + unsafe fn as_bytes_mut(self) -> &'a mut [u8] { + slice::from_raw_parts_mut(self.data(), self.len()) + } +} + impl<'py> TryFrom<&'py PyAny> for &'py PyByteArray { type Error = crate::PyErr; diff --git a/src/types/bytes.rs b/src/types/bytes.rs index 4b5b0fb3..5b550dd3 100644 --- a/src/types/bytes.rs +++ b/src/types/bytes.rs @@ -1,4 +1,4 @@ -use crate::instance::Py2; +use crate::instance::{Py2, Py2Borrowed}; use crate::{ffi, FromPyObject, IntoPy, Py, PyAny, PyResult, Python, ToPyObject}; use std::borrow::Cow; use std::ops::Index; @@ -89,9 +89,7 @@ impl PyBytes { /// Gets the Python string as a byte slice. #[inline] pub fn as_bytes(&self) -> &[u8] { - let slice = Py2::borrowed_from_gil_ref(&self).as_bytes(); - // SAFETY: &self guarantees the reference is alive long enough - unsafe { std::slice::from_raw_parts(slice.as_ptr(), slice.len()) } + Py2Borrowed::from_gil_ref(self).as_bytes() } } @@ -109,6 +107,14 @@ pub(crate) trait PyBytesMethods<'py> { impl<'py> PyBytesMethods<'py> for Py2<'py, PyBytes> { #[inline] fn as_bytes(&self) -> &[u8] { + Py2Borrowed::from(self).as_bytes() + } +} + +impl<'a> Py2Borrowed<'a, '_, PyBytes> { + /// Gets the Python string as a byte slice. + #[allow(clippy::wrong_self_convention)] + fn as_bytes(self) -> &'a [u8] { unsafe { let buffer = ffi::PyBytes_AsString(self.as_ptr()) as *const u8; let length = ffi::PyBytes_Size(self.as_ptr()) as usize; @@ -123,10 +129,7 @@ impl Py { /// immutable, the result may be used for as long as the reference to /// `self` is held, including when the GIL is released. pub fn as_bytes<'a>(&'a self, py: Python<'_>) -> &'a [u8] { - let slice = self.attach(py).as_bytes(); - // SAFETY: it is safe to access the immutable slice as long as the - // reference to `self` is held. - unsafe { std::slice::from_raw_parts(slice.as_ptr(), slice.len()) } + self.attach_borrow(py).as_bytes() } }