Python Iterator for Generic Rust Struct

[StuartHadfield]

Hello!

Me again. You probably know the story by now, I have a rust struct wrapping a generic vector for use in a python pipeline. The idea is to speed up array handling in Python.

It looks like this (big thanks to contributors here for getting me this far):

#[pyclass]
pub struct Wrapper {
    inner: WrapperInner,
}

#[pymethods]
impl Wrapper {
    #[new]
    fn new(members: &PyAny) -> PyResult<Self> {
        let inner = if let Ok(members) = Vec::<i64>::extract(members) {
            WrapperInner::Int(Core { members })
        } else if let Ok(members) = Vec::<String>::extract(members) {
            WrapperInner::Str(Core { members })
     // etc...
}

Then, I have a core.rs containing the Rust Vector processing side:

pub struct Core<T> {
    pub members: Vec<T>,
}

impl<T> Core<T>
where // Traits I need
{
// fns
}

Now, I'm trying to implement iteration on this struct (e.g. iter([1,2,3]).next() kinda thing in python). I managed this with a strictly typed Vec as follows:

impl Core {
    // fns ...
    fn __iter__(slf: PyRef<Self>) -> PyResult<Py<Iter>> {
        let iter = Iter {
            inner: slf.members.clone().into_iter(),
        };
        Py::new(slf.py(), iter)
    }
}

#[pyclass]
struct Iter {
    inner: std::vec::IntoIter<String>,
}

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

    fn __next__(mut slf: PyRefMut<Self>) -> Option<String> {
        slf.inner.next()
    }
}

Which was all fine and good - I was always sure of getting a String iterator out which made things easy. I tried to adapt this as follows, starting with ONLY int or string cases (I can extrapolate from there):

// in wrapper.rs - I don't want to pollute core.rs with pyo3 context
impl Wrapper {
    // fns...

    fn iter(slf: PyRef<Self>) -> Result<Py<StrOrIntIter>, PyErr> {
        let iter: StrOrIntIter = match &slf.inner {
            WrapperInner::Str(core) => {
                StrOrIntIter::Str(IterStr {
                    inner: core.members.clone().into_iter(),
                })
            },
            WrapperInner::Int(core) => {
                StrOrIntIter::Int(IterInt {
                    inner: core.members.clone().into_iter(),
                })
            },
        };

        Py::new(slf.py(), iter)
}

#[pyclass]
enum StrOrIntIter {
    Str(IterStr),
    Int(IterInt),
}

[pyclass]
pub struct IterInt {
    inner: std::vec::IntoIter<i64>,
}

#[pyclass]
pub struct IterStr {
    inner: std::vec::IntoIter<String>,
}

#[pymethods]
impl IterInt {
    fn __iter__(slf: PyRef<Self>) -> PyRef<Self> {
        slf
    }

    fn __next__(mut slf: PyRefMut<Self>) -> Option<i64> {
        slf.inner.next()
    }
}

#[pymethods]
impl IterStr {
    fn __iter__(slf: PyRef<Self>) -> PyRef<Self> {
        slf
    }

    fn __next__(mut slf: PyRefMut<Self>) -> Option<String> {
        slf.inner.next()
    }
}

Of course, this won't work, because Pyo3 has no idea how to turn a StrOrIntIter into a python iterator (though it would understand how to return a IterInt or a IterStr)

My error message looks something like:

error[E0277]: the trait bound `StrOrIntIter: PyClass` is not satisfied                                                                                                     
   --> ...
    |
201 |         Py::new(slf.py(), iter)
    |         -------           ^^^^ the trait `PyClass` is not implemented for `StrOrIntIter`
    |         |
    |         required by a bound introduced by this call
    |
    = note: required because of the requirements on the impl of `From<StrOrIntIter>` for `PyClassInitializer<StrOrIntIter>`
    = note: required because of the requirements on the impl of `Into<PyClassInitializer<StrOrIntIter>>` for `StrOrIntIter`
note: required by a bound in `pyo3::Py::<T>::new`
   --> ...
    |
254 |     pub fn new(py: Python<'_>, value: impl Into<PyClassInitializer<T>>) -> PyResult<Py<T>> {
    |                                            ^^^^^^^^^^^^^^^^^^^^^^^^^^^ required by this bound in `pyo3::Py::<T>::new`

Now, I know exactly what this is telling me. I have to implement Into and From on StrOrInter, but I have nooo idea what this implementation signature looks like in the context of Pyo3. I've guessed at:

impl Into<PyClassInitializer<T>> for StrOrIntIter {
    fn into(self, py: Python<'_>) => PyObject {
        match self {
            StrOrIntIter::Str(iter) -> Py::new(PyRef::new(py, self).py(), iter),
            StrOrIntIter::Int(iter) -> Py::new(PyRef::new(py, self).py(), iter),
        }
    }
}

But this has got me so many errors I don't know where to start. I can't find the documentation on this - any guidance would be appreciated!

[birkenfeld]

Implemeting Into<PyClassInitializer> is not going to help here. You could implement PyClass for StrOrIntIter, forwarding the __iter__ and __next__ methods to the inner types.

However, if StrOrIntIter is only needed for the benefit of Python code, I think it would be easier to have the iter return Py<PyAny>, without a wrapper type, somewhat like this:

// in wrapper.rs - I don't want to pollute core.rs with pyo3 context
impl Wrapper {
    // fns...

    fn iter(&self, py: Python<'_>) -> Result<Py<PyAny>, PyErr> {
        match &self.inner {
            WrapperInner::Str(core) => {
                IterStr {
                    inner: core.members.clone().into_iter(),
                }.into_py(py)
            },
            WrapperInner::Int(core) => {
                IterInt {
                    inner: core.members.clone().into_iter(),
                }.into_py(py)
            },
        }
}

Cloning the inner vector is not very efficient, by the way. Depending on the circumstances, it would probably be better to implement the IterInt/Str types by keeping a reference to the wrapper and a running index...

[StuartHadfield]

Hah, I feel really dumb now. That simplifies a lot of my code - I didn't realise you could return PyAny by calling into_py. Thanks Georg! This is the right answer for me.

Out of interest, what do you mean by "keeping a reference to the wrapper and a running index"?

[birkenfeld]

The Iterator type would look like this:

struct Iterator {
    wrapper: Py<Wrapper>,
    index: usize
}

and the __next__ would then retrieve and return wrapper[index] (plus check for end of iteration, increment the index, etc.).

[mejrs]

You can't use #[pyclass] on an enum with fields.

[StuartHadfield]

Thanks Bruno! I must've got confused reading the documentation 😞