Search code examples
rustpyo3

Rust, PyO3, and return a "simple" internal iterator


Older answers use the outdated #[pyprotocol] workflow for magic methods. I think PyO3 has matured a bit, so I'm asking this again.

I have a simple struct

#[pyclass]
struct Outer {
  inner: Vec<i32>
}

In the PyO3 documentation, if I want to make Outer iterable, I need to implement __iter__, which needs to return something that implements both __iter__ and __next. So in the documentation, they provide a

#[pyclass]
Iter {
  iter: std::Vec::IntoIter<i32>
}

#[pymethods]
impl Iter {
    fn __iter__(slf: PyRef<'_, Self>) -> PyRef<'_, Self> {
        slf
    }

    fn __next__(mut slf: PyRefMut<'_, Self>) -> Option<bool> {
        slf.inner.next()
    }
}

and then my Outer will have

pub fn __iter__(slf: PyRef<'_, Self>) -> PyResult<Py<Iter>> {
    let iter = Iter {
        inner: slf.config.clone().into_iter(),
    };
    Py::new(slf.py(), iter)
}

Is there no shorter way to do this? For example, Vec<i32>, thanks to PyO3, will implement the necessary traits for automatic conversion into a python list. Aren't there similar shortcuts so that I wouldn't have to implement that whole Iter struct and could just do something like

fn __iter__(...) {
  self.inner.iter().cloned()
}

and maybe wrap it in the correct PyResult or whatever things? I'm fine doing it with the extra struct, it just feels an extra step that might not be needed.


Solution

  • You can return an existing Python iterator (for example, a list iterator) from __iter__(). This will be easier because you can use PyO3 existing machinery for converting Rust collections to Python collections, but slower because you need to copy the whole Rust collection to a Python collection:

    use pyo3::types::PyIterator;
    
    #[pymethods]
    impl Outer {
        pub fn __iter__<'py>(&self, py: Python<'py>) -> PyResult<&'py PyIterator> {
            PyIterator::from_object(self.inner.to_object(py).into_ref(py))
        }
    }