How can I access a Rust Iterator from Python using PyO3?

Question

I'm quite new with Rust, and my first 'serious' project has involved writing a Python wrapper for a small Rust library using PyO3. This has mostly been quite painless, but I'm struggling to work out how to expose lazy iterators over Rust Vec s to Python code.

So far, I have been collecting the values produced by the iterator and returning a list, which obviously isn't the best solution. Here's some code which illustrates my problem:

use pyo3::prelude::*;

// The Rust Iterator, from the library I'm wrapping.
pub struct RustIterator<'a> {
    position: usize,
    view: &'a Vec<isize>
}

impl<'a> Iterator for RustIterator<'a> {
    type Item = &'a isize;

    fn next(&mut self) -> Option<Self::Item> {
        let result = self.view.get(self.position);
        if let Some(_) = result { self.position += 1 };
        result
    }
}

// The Rust struct, from the library I'm wrapping.
struct RustStruct {
    v: Vec<isize>
}

impl RustStruct {
    fn iter(&self) -> RustIterator {
        RustIterator{ position: 0, view: &self.v }
    }
}

// The Python wrapper class, which exposes the 
// functions of RustStruct in a Python-friendly way.
#[pyclass]
struct PyClass {
    rust_struct: RustStruct,
}

#[pymethods]
impl PyClass {
    #[new]
    fn new(v: Vec<isize>) -> Self {
        let rust_struct = RustStruct { v };
        Self{ rust_struct }
    }

    // This is what I'm doing so far, which works
    // but doesn't iterate lazily.
    fn iter(&self) -> Vec<isize> {
        let mut output_v = Vec::new();
        for item in self.rust_struct.iter() {
            output_v.push(*item);
        }
        output_v
    }
}

I've tried to wrap the RustIterator class with a Python wrapper, but I can't use PyO3's #[pyclass] proc. macro with lifetime parameters. I looked into pyo3::types::PyIterator but this looks like a way to access a Python iterator from Rust rather than the other way around.

How can I access a lazy iterator over RustStruct.v in Python? It's safe to assume that the type contained in the Vec always derives Copy and Clone , and answers which require some code on the Python end are okay (but less ideal).

Answer 1

You can make your RustIterator a pyclass and then implement the proper trait ( PyIterProtocol ) using the rust iter itself.

Not tested, but something like:

#[pyclass]
pub struct RustIterator<'a> {
    position: usize,
    view: &'a Vec<isize>
}

impl<'a> Iterator for RustIterator<'a> {
    type Item = &'a isize;

    fn next(&mut self) -> Option<Self::Item> {
        let result = self.view.get(self.position);
        if let Some(_) = result { self.position += 1 };
        result
    }
}

#[pyproto]
impl PyIterProtocol for Iter {
    fn __next__(mut slf: PyRefMut<Self>) -> IterNextOutput<usize, &'static str> {
        match self.next() {
            Some(value) => IterNextOutput::Yield(value),
            None => IterNextOutput::Return("Ended")
        }
    }
}