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