Mercurial > public > mercurial-scm > hg
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 {