use Py2Borrowed to make PyBytesMethods slightly nicer

Author: davidhewitt

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<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)
+    }
 }

src/instance.rs

@@ -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).
     ///

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,13 @@ 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.
+    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 +128,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()
     }
 }