Question: How can I make a pymethod return a generator?

[thesketh]

Hi all,

I'm writing a Python wrapper for a Rust lib, and PyO3 has been great so far! It's been a pleasure to use, and I was so impressed by the fact that PyPy was supported with literally no effort on my part.

Some of the Rust methods I've wrapped would normally return Iterators over Vecs in the data, which mostly look a bit like this:

pub struct SomeIterator<'a> {
    c: usize,
    v: &'a Vec<(usize, usize)>  // reference to Vec owned by RustStruct
}

I can't wrap these using #[pyclass] because they contain a lifetime annotation. I came up with something like this instead:

struct RustStruct {
    v1: Vec<isize>,
    v2: Vec<(usize, usize)>
}

#[pyclass]
struct PyClass {
    rust_struct: RustStruct,
    c1: usize,
    c2: usize
}

#[pymethods]
impl PyClass {
    #[new]
    fn new(v1: Vec<isize>, v2: Vec<(usize, usize)>) -> Self {
        Self{ rust_struct: RustStruct { v1, v2 }, c1: 0, c2: 0 }
    }

    fn next_v1(&mut self) -> PyResult<Option<isize>> {
        match self.rust_struct.v1.get(self.c1) {
            Some(v) => {
                self.c1 += 1;
                Ok(Some(*v))
            },
            None => {
                self.c1 = 0;
                Ok(None)
            }
        }
    }

    fn next_v2(&mut self) -> PyResult<Option<(usize, usize)>> {
        match self.rust_struct.v2.get(self.c2) {
            Some((v1, v2)) => {
                self.c2 += 1;
                Ok(Some((*v1, *v2)))
            },
            None => {
                self.c2 = 0;
                Ok(None)
            }
        }
    }
}

I was hoping to be able to subclass PyClass like so, and define some functions which return a Python generator:

from some_module import PyClass

class Class(PyClass):
    def __new__(cls, l1, l2):
        return PyClass.__new__(cls, l1, l2)

    @staticmethod
    def _gen_from_next(next_fn):
        while True:
            val = self.next_fn()
            if val is None:
                return None
            yield val

    def iter_v1(self):
        return self._gen_from_next(self.next_v1)

    def iter_v2(self):
        return self._gen_from_next(self.next_v2)

But it seems that I can't properly subclass PyO3 #[pyclass] structs: Class.__new__ actually returns an instance of PyClass so I'm unable to use iter_v1 and iter_v2.

I'm trying to avoid cloning the whole Vec at once: if I was going to clone the Vecs anyway, I could add them as members of PyClass and use #[pyo3(get)] to access them as lists in Python.

Any ideas?

[althonos]

AFAIK you can't have generators on the Rust side, but you can return an iterator being declared in Rust as a #[pyclass] as well.

Btw, have you been declaring you class with #[pyclass(subclass)]? (Just to know, I don't know if it's a fix or not since it keeps changing but this may be needed for you to be able to subclass your Rust class)

[thesketh]

RE: the first point, I've only been able to get this to work when the iterator owns the data it is iterating over. If the class had a single function which returned an iterator, I could implement __iter__ and __next__ without much issue, but in this case I have a struct with three functions which should return iterators. If I try to create a #[pyclass] with a lifetime parameter for the reference I see error: #[pyclass] cannot have generic parameters.

I am using #[pyclass(subclass)], but ParentClass.__new__ seems to ignore the first argument and always returns an instance of the parent class instead. The only thing it seems I can do with subclassing is create new constructors

[davidhewitt]

ParentClass.new seems to ignore the first argument and always returns an instance of the parent class instead

That sounds like a bug. I've opened a new issue #947 to track it when someone has a chance to investigate.

@thesketh , regarding the ownership of the data, that's correct. There's no way to express a Rust lifetime safely in a #[pyclass].

To avoid having to copy the vector, you could consider using a shared pointer like Arc or Py in the iterator. For example, your iterator implementation could be:

#[pyclass]
pub struct SomePyIterator<'a> {
    c: usize,
    v: Py<RustStruct> // reference to RustStruct instantiated in Python
}

[thesketh]

Thanks for opening that, I'll keep an eye on it.

And thanks for the pointer (no pun intended), I'm quite new to using Rust and I didn't realise that using the reference counted types was the right solution to this. That answers my question!