Mercurial > public > mercurial-scm > hg

comparison rust/pyo3-sharedref/src/lib.rs @ 52860:8f6d25439bdc

Find changesets by keywords (author, files, the commit message), revision number or hash, or revset expression.
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 189491cea922
children
comparison
equal deleted inserted replaced
52859:9f083ff3c96c 52860:8f6d25439bdc
744 impl<'a, T: ?Sized> DerefMut for SharedByPyObjectRefMut<'a, T> { 744 impl<'a, T: ?Sized> DerefMut for SharedByPyObjectRefMut<'a, T> {
745 fn deref_mut(&mut self) -> &mut T { 745 fn deref_mut(&mut self) -> &mut T {
746 self.data 746 self.data
747 } 747 }
748 } 748 }
749
750 /// Defines a Python iterator over a Rust iterator.
751 ///
752 /// TODO: this is a bit awkward to use, and a better (more complicated)
753 /// procedural macro would simplify the interface a lot.
754 ///
755 /// # Parameters
756 ///
757 /// * `$name` is the identifier to give to the resulting Rust struct.
758 /// * `$success_type` is the resulting Python object. It can be a bultin type,
759 /// (e.g., `PyBytes`), or any `PyClass`.
760 /// * `$owner_type` is the type owning the data
761 /// * `$owner_attr` is the name of the shareable attribute in `$owner_type`
762 /// * `$shared_type` is the type wrapped in `SharedByPyObject`, typically
763 /// `SomeIter<'static>`
764 /// * `$iter_func` is a function to obtain the Rust iterator from the content
765 /// of the shareable attribute. It can be a closure.
766 /// * `$result_func` is a function for converting items returned by the Rust
767 /// iterator into `PyResult<Option<Py<$success_type>`.
768 ///
769 /// # Safety
770 ///
771 /// `$success_func` may take a reference, whose lifetime may be articial.
772 /// Do not copy it out of the function call (this would be possible only
773 /// with inner mutability).
774 ///
775 /// # Example
776 ///
777 /// The iterator example in [`PyShareable`] can be rewritten as
778 ///
779 /// ```
780 /// use pyo3::prelude::*;
781 /// use pyo3_sharedref::*;
782 ///
783 /// use pyo3::types::{PyTuple, PyInt};
784 /// use std::collections::{hash_set::Iter as IterHashSet, HashSet};
785 /// use std::vec::Vec;
786 ///
787 /// #[pyclass(sequence)]
788 /// struct Set {
789 /// rust_set: PyShareable<HashSet<i32>>,
790 /// }
791 ///
792 /// #[pymethods]
793 /// impl Set {
794 /// #[new]
795 /// fn new(values: &Bound<'_, PyTuple>) -> PyResult<Self> {
796 /// let as_vec = values.extract::<Vec<i32>>()?;
797 /// let s: HashSet<_> = as_vec.iter().copied().collect();
798 /// Ok(Self { rust_set: s.into() })
799 /// }
800 ///
801 /// fn __iter__(slf: &Bound<'_, Self>) -> PyResult<SetIterator> {
802 /// SetIterator::new(slf)
803 /// }
804 /// }
805 ///
806 /// py_shared_iterator!(
807 /// SetIterator,
808 /// PyInt,
809 /// Set,
810 /// rust_set,
811 /// IterHashSet<'static, i32>,
812 /// |hash_set| hash_set.iter(),
813 /// it_next_result
814 /// );
815 ///
816 /// fn it_next_result(py: Python, res: &i32) -> PyResult<Option<Py<PyInt>>> {
817 /// Ok(Some((*res).into_pyobject(py)?.unbind()))
818 /// }
819 /// ```
820 ///
821 /// In the example above, `$result_func` is fairly trivial, and can be replaced
822 /// by a closure, but things can get more complicated if the Rust
823 /// iterator itself returns `Result<T, E>` with `T` not implementing
824 /// `IntoPyObject` and `E` needing to be converted.
825 /// Also the closure variant is fairly obscure:
826 ///
827 /// ```ignore
828 /// py_shared_iterator!(
829 /// SetIterator,
830 /// PyInt,
831 /// Set,
832 /// rust_set,
833 /// IterHashSet<'static, i32>,
834 /// |hash_set| hash_set.iter(),
835 /// (|py, i: &i32| Ok(Some((*i).into_pyobject(py)?.unbind())))
836 /// )
837 /// ```
838 #[macro_export]
839 macro_rules! py_shared_iterator {
840 (
841 $name: ident,
842 $success_type: ty,
843 $owner_type: ident,
844 $owner_attr: ident,
845 $shared_type: ty,
846 $iter_func: expr,
847 $result_func: expr
848 ) => {
849 #[pyclass]
850 pub struct $name {
851 inner: pyo3_sharedref::SharedByPyObject<$shared_type>,
852 }
853
854 #[pymethods]
855 impl $name {
856 #[new]
857 fn new(owner: &Bound<'_, $owner_type>) -> PyResult<Self> {
858 let inner = &owner.borrow().$owner_attr;
859 // Safety: the data is indeed owned by `owner`
860 let shared_iter =
861 unsafe { inner.share_map(owner, $iter_func) };
862 Ok(Self { inner: shared_iter })
863 }
864
865 fn __iter__(slf: PyRef<'_, Self>) -> PyRef<'_, Self> {
866 slf
867 }
868
869 fn __next__(
870 mut slf: PyRefMut<'_, Self>,
871 ) -> PyResult<Option<Py<$success_type>>> {
872 let py = slf.py();
873 let shared = &mut slf.inner;
874 // Safety: we do not leak references derived from the internal
875 // 'static reference.
876 let mut inner = unsafe { shared.try_borrow_mut(py) }?;
877 match inner.next() {
878 None => Ok(None),
879 Some(res) => $result_func(py, res),
880 }
881 }
882 }
883 };
884 }