Mercurial > public > mercurial-scm > hg
diff rust/hg-pyo3/src/util.rs @ 52836:9749a97d3cfb
rust-pyo3: introduce utils to get the pyo3-wrapped index from Python
This is the equivalent of what we did in `convert_cpython`, but with fewer
steps, since we don't have to do the convertion/checking from cpython to pyo3
anymore.
This will be used in the next changeset when plugging everything in.
author | Rapha?l Gom?s <rgomes@octobus.net> |
---|---|
date | Sun, 05 Jan 2025 23:33:24 +0100 |
parents | dd3a2948804f |
children |
line wrap: on
line diff
--- a/rust/hg-pyo3/src/util.rs Sun Jan 05 23:21:56 2025 +0100 +++ b/rust/hg-pyo3/src/util.rs Sun Jan 05 23:33:24 2025 +0100 @@ -1,11 +1,15 @@ use hg::errors::HgError; +use hg::revlog::index::Index as CoreIndex; use hg::revlog::inner_revlog::RevisionBuffer; use pyo3::buffer::{Element, PyBuffer}; use pyo3::exceptions::PyValueError; use pyo3::prelude::*; use pyo3::types::{PyBytes, PyDict}; +use pyo3_sharedref::SharedByPyObject; use stable_deref_trait::StableDeref; +use crate::revlog::{InnerRevlog, PySharedIndex}; + /// Create the module, with `__package__` given from parent /// /// According to PyO3 documentation, which links to @@ -33,6 +37,70 @@ Ok(m) } +/// Retrieve the shared index wrapper (which contains the core index) +/// from the Python index proxy. +pub fn py_rust_index_to_graph( + index_proxy: &Bound<'_, PyAny>, +) -> PyResult<SharedByPyObject<PySharedIndex>> { + let py_irl = index_proxy.getattr("inner")?; + let py_irl_ref = py_irl.downcast::<InnerRevlog>()?.borrow(); + let shareable_irl = &py_irl_ref.irl; + + // Safety: the owner is the actual one and we do not leak any + // internal reference. + let index = + unsafe { shareable_irl.share_map(&py_irl, |irl| (&irl.index).into()) }; + Ok(index) +} + +/// Error propagation for an [`SharedByPyObject`] wrapping a [`Result`] +/// +/// It would be nice for [`SharedByPyObject`] to provide this directly as +/// a variant of the `map` method with a signature such as: +/// +/// ``` +/// unsafe fn map_or_err(&self, +/// py: Python, +/// f: impl FnOnce(T) -> Result(U, E), +/// convert_err: impl FnOnce(E) -> PyErr) +/// ``` +/// +/// This would spare users of the `pyo3` crate the additional `unsafe` deref +/// to inspect the error and return it outside `SharedByPyObject`, and the +/// subsequent unwrapping that this function performs. +pub(crate) fn py_shared_or_map_err<T, E: std::fmt::Debug + Copy>( + py: Python, + leaked: SharedByPyObject<Result<T, E>>, + convert_err: impl FnOnce(E) -> PyErr, +) -> PyResult<SharedByPyObject<T>> { + // Safety: we don't leak the "faked" reference out of `SharedByPyObject` + if let Err(e) = *unsafe { leaked.try_borrow(py)? } { + return Err(convert_err(e)); + } + // Safety: we don't leak the "faked" reference out of `SharedByPyObject` + Ok(unsafe { + leaked.map(py, |res| { + res.expect("Error case should have already be treated") + }) + }) +} + +/// Full extraction of the proxy index object as received in PyO3 to a +/// [`CoreIndex`] reference. +/// +/// # Safety +/// +/// The invariants to maintain are those of the underlying +/// [`SharedByPyObject::try_borrow`]: the caller must not leak the inner +/// reference. +pub(crate) unsafe fn proxy_index_extract<'py>( + index_proxy: &Bound<'py, PyAny>, +) -> PyResult<&'py CoreIndex> { + let py_shared = py_rust_index_to_graph(index_proxy)?; + let py_shared = &*unsafe { py_shared.try_borrow(index_proxy.py())? }; + Ok(unsafe { py_shared.static_inner() }) +} + /// Type shortcut for the kind of bytes slice trait objects that are used in /// particular for mmap data type BoxedBytesSlice =