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

42 Views Asked by At

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.

1

There are 1 best solutions below

1
Chayim Friedman On BEST ANSWER

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))
    }
}