use Py2Borrowed to make PyBytesMethods slightly nicer

This commit is contained in:
David Hewitt 2023-12-14 19:07:47 +00:00
parent c5dce0172b
commit 35f7f1a78c
4 changed files with 171 additions and 19 deletions

View File

@ -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<Py2<'_, PyAny>>;
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<Py2Borrowed<'a, '_, PyAny>>;
/// 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<Py2Borrowed<'a, '_, PyAny>>;
/// 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<'a, '_, PyAny>> {
Py2Borrowed::from_ptr_or_err(py, self)
}
#[inline]
unsafe fn assume_borrowed_or_opt<'a>(
self,
py: Python<'_>,
) -> Option<Py2Borrowed<'a, '_, PyAny>> {
Py2Borrowed::from_ptr_or_opt(py, self)
}
#[inline]
unsafe fn assume_borrowed<'a>(self, py: Python<'_>) -> Py2Borrowed<'a, '_, PyAny> {
Py2Borrowed::from_ptr(py, self)
}
}

View File

@ -217,6 +217,96 @@ unsafe impl<T> 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<ffi::PyObject>,
PhantomData<&'a Py<T>>,
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<Self> {
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<Self> {
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<T> 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<ffi::PyObject>
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<T> Py<T> {
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).
///

View File

@ -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<u8> {
@ -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;

View File

@ -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<PyBytes> {
/// 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()
}
}