diff rust/hg-pyo3/src/util.rs @ 52830:dd3a2948804f

rust-pyo3: add `PyBytesDeref` util This mirrors the one in hg-cpython, explanations inline.
author Rapha?l Gom?s <rgomes@octobus.net>
date Fri, 03 Jan 2025 15:44:29 +0100
parents 4f41a8acf350
children 9749a97d3cfb
line wrap: on
line diff
--- a/rust/hg-pyo3/src/util.rs	Fri Jan 03 15:00:50 2025 +0100
+++ b/rust/hg-pyo3/src/util.rs	Fri Jan 03 15:44:29 2025 +0100
@@ -4,6 +4,7 @@
 use pyo3::exceptions::PyValueError;
 use pyo3::prelude::*;
 use pyo3::types::{PyBytes, PyDict};
+use stable_deref_trait::StableDeref;
 
 /// Create the module, with `__package__` given from parent
 ///
@@ -114,6 +115,62 @@
     }
 }
 
+/// Safe abstraction over a `PyBytes` together with the `&[u8]` slice
+/// that borrows it. Implements `Deref<Target = [u8]>`.
+///
+/// Calling `PyBytes::data` requires a GIL marker but we want to access the
+/// data in a thread that (ideally) does not need to acquire the GIL.
+/// This type allows separating the call an the use.
+///
+/// It also enables using a (wrapped) `PyBytes` in GIL-unaware generic code.
+pub struct PyBytesDeref {
+    #[allow(unused)]
+    keep_alive: Py<PyBytes>,
+
+    /// Borrows the buffer inside `self.keep_alive`,
+    /// but the borrow-checker cannot express self-referential structs.
+    data: *const [u8],
+}
+
+impl PyBytesDeref {
+    pub fn new(py: Python, bytes: Py<PyBytes>) -> Self {
+        Self {
+            data: bytes.as_bytes(py),
+            keep_alive: bytes,
+        }
+    }
+
+    #[allow(unused)]
+    pub fn unwrap(self) -> Py<PyBytes> {
+        self.keep_alive
+    }
+}
+
+impl std::ops::Deref for PyBytesDeref {
+    type Target = [u8];
+
+    fn deref(&self) -> &[u8] {
+        // Safety: the raw pointer is valid as long as the PyBytes is still
+        // alive, and the returned slice borrows `self`.
+        unsafe { &*self.data }
+    }
+}
+
+unsafe impl StableDeref for PyBytesDeref {}
+
+fn require_send<T: Send>() {}
+
+#[allow(unused)]
+fn static_assert_pybytes_is_send() {
+    #[allow(clippy::no_effect)]
+    require_send::<Py<PyBytes>>;
+}
+
+// Safety: `[Py<PyBytes>]` is Send. Raw pointers are not by default,
+// but here sending one to another thread is fine since we ensure it stays
+// valid.
+unsafe impl Send for PyBytesDeref {}
+
 /// Wrapper around a Python-provided buffer into which the revision contents
 /// will be written. Done for speed in order to save a large allocation + copy.
 struct PyRevisionBuffer {