rust/pyo3-sharedref/src/lib.rs
changeset 52860 8f6d25439bdc
parent 52843 189491cea922
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 }