Mercurial > public > mercurial-scm > hg-stable
changeset 52872:8f6d25439bdc
rust-pyo3-sharedref: macro to define Python iterators
This convenience will help reducing the boilerplate for one of
the most common use cases of the `pyo3_sharedref` crate.
It is an improved adaptation of a similar macro that was in `hg-cpython`
(but not in `rust-cpython`). The improvement is that it takes care of
the constructor, sparing the `share_map` to the caller.
author | Georges Racinet <georges.racinet@cloudcrane.io> |
---|---|
date | Tue, 04 Feb 2025 11:55:14 +0100 |
parents | 9f083ff3c96c |
children | c5773445d350 |
files | rust/pyo3-sharedref/src/lib.rs |
diffstat | 1 files changed, 136 insertions(+), 0 deletions(-) [+] |
line wrap: on
line diff
--- a/rust/pyo3-sharedref/src/lib.rs Wed Jan 29 18:26:10 2025 +0100 +++ b/rust/pyo3-sharedref/src/lib.rs Tue Feb 04 11:55:14 2025 +0100 @@ -746,3 +746,139 @@ self.data } } + +/// Defines a Python iterator over a Rust iterator. +/// +/// TODO: this is a bit awkward to use, and a better (more complicated) +/// procedural macro would simplify the interface a lot. +/// +/// # Parameters +/// +/// * `$name` is the identifier to give to the resulting Rust struct. +/// * `$success_type` is the resulting Python object. It can be a bultin type, +/// (e.g., `PyBytes`), or any `PyClass`. +/// * `$owner_type` is the type owning the data +/// * `$owner_attr` is the name of the shareable attribute in `$owner_type` +/// * `$shared_type` is the type wrapped in `SharedByPyObject`, typically +/// `SomeIter<'static>` +/// * `$iter_func` is a function to obtain the Rust iterator from the content +/// of the shareable attribute. It can be a closure. +/// * `$result_func` is a function for converting items returned by the Rust +/// iterator into `PyResult<Option<Py<$success_type>`. +/// +/// # Safety +/// +/// `$success_func` may take a reference, whose lifetime may be articial. +/// Do not copy it out of the function call (this would be possible only +/// with inner mutability). +/// +/// # Example +/// +/// The iterator example in [`PyShareable`] can be rewritten as +/// +/// ``` +/// use pyo3::prelude::*; +/// use pyo3_sharedref::*; +/// +/// use pyo3::types::{PyTuple, PyInt}; +/// use std::collections::{hash_set::Iter as IterHashSet, HashSet}; +/// use std::vec::Vec; +/// +/// #[pyclass(sequence)] +/// struct Set { +/// rust_set: PyShareable<HashSet<i32>>, +/// } +/// +/// #[pymethods] +/// impl Set { +/// #[new] +/// fn new(values: &Bound<'_, PyTuple>) -> PyResult<Self> { +/// let as_vec = values.extract::<Vec<i32>>()?; +/// let s: HashSet<_> = as_vec.iter().copied().collect(); +/// Ok(Self { rust_set: s.into() }) +/// } +/// +/// fn __iter__(slf: &Bound<'_, Self>) -> PyResult<SetIterator> { +/// SetIterator::new(slf) +/// } +/// } +/// +/// py_shared_iterator!( +/// SetIterator, +/// PyInt, +/// Set, +/// rust_set, +/// IterHashSet<'static, i32>, +/// |hash_set| hash_set.iter(), +/// it_next_result +/// ); +/// +/// fn it_next_result(py: Python, res: &i32) -> PyResult<Option<Py<PyInt>>> { +/// Ok(Some((*res).into_pyobject(py)?.unbind())) +/// } +/// ``` +/// +/// In the example above, `$result_func` is fairly trivial, and can be replaced +/// by a closure, but things can get more complicated if the Rust +/// iterator itself returns `Result<T, E>` with `T` not implementing +/// `IntoPyObject` and `E` needing to be converted. +/// Also the closure variant is fairly obscure: +/// +/// ```ignore +/// py_shared_iterator!( +/// SetIterator, +/// PyInt, +/// Set, +/// rust_set, +/// IterHashSet<'static, i32>, +/// |hash_set| hash_set.iter(), +/// (|py, i: &i32| Ok(Some((*i).into_pyobject(py)?.unbind()))) +/// ) +/// ``` +#[macro_export] +macro_rules! py_shared_iterator { + ( + $name: ident, + $success_type: ty, + $owner_type: ident, + $owner_attr: ident, + $shared_type: ty, + $iter_func: expr, + $result_func: expr + ) => { + #[pyclass] + pub struct $name { + inner: pyo3_sharedref::SharedByPyObject<$shared_type>, + } + + #[pymethods] + impl $name { + #[new] + fn new(owner: &Bound<'_, $owner_type>) -> PyResult<Self> { + let inner = &owner.borrow().$owner_attr; + // Safety: the data is indeed owned by `owner` + let shared_iter = + unsafe { inner.share_map(owner, $iter_func) }; + Ok(Self { inner: shared_iter }) + } + + fn __iter__(slf: PyRef<'_, Self>) -> PyRef<'_, Self> { + slf + } + + fn __next__( + mut slf: PyRefMut<'_, Self>, + ) -> PyResult<Option<Py<$success_type>>> { + let py = slf.py(); + let shared = &mut slf.inner; + // Safety: we do not leak references derived from the internal + // 'static reference. + let mut inner = unsafe { shared.try_borrow_mut(py) }?; + match inner.next() { + None => Ok(None), + Some(res) => $result_func(py, res), + } + } + } + }; +}