How to implement __iter__ when yielding a #[pyclass] type?

[aeshirey]

I'm trying to implement __iter__ and __next__ for my own types, but I'm struggling to figure it out. I can make it work if the values I'm yielding are Python types (eg, PyString), but not when I'm yielding my own type. Here's a minimal example that is probably very close, but no cigar:

use pyo3::{
    Bound, Py, PyResult, Python, pyclass, pymethods, pymodule,
    types::{PyModule, PyModuleMethods},
};

#[pyclass]
struct Item {
    value: String,
}

#[pymethods]
impl Item {
    #[getter]
    fn value(&self) -> &str {
        &self.value
    }
}

#[pyclass]
struct Container {
    items: Vec<Item>,
}

#[pymethods]
impl Container {
    #[new]
    fn new() -> Self {
        Container { items: Vec::new() }
    }

    fn __iter__(&self, py: Python) -> ContainerIterator {
        ContainerIterator {
            index: 0,
            container: todo!(), 
        }
    }
}

#[pyclass]
struct ContainerIterator {
    index: usize,
    container: Py<Container>, // not sure if this is the right type
}

impl ContainerIterator {
    fn __next__(&mut self, py: Python) -> PyResult<Option<Item>> {
        // this isn't right:
        if self.index < self.container.items.len() {
            let item = self.container.items[self.index].clone(); // I don't want to clone the underlying value
            self.index += 1;
            Ok(Some(item))
        } else {
            Ok(None)
        }
    }
}

#[pymodule]
fn pyiter(m: &Bound<'_, PyModule>) -> PyResult<()> {
    m.add_class::<Container>()?;
    m.add_class::<Item>()?;

    Ok(())
}

What am I missing here? Presumably, the ContainerIterator needs to have some sort of a reference to the Container, and its __next__ method should yield some sort of a reference back to Python?

[aeshirey]

I'll add that, from the Python side, use of this would be something like:

c = Container() # how this gets populated with values will be addressed elsewhere

for (i, item) in enumerate(c):
    # note that item is of type Item, and it has a .value getter:
    print(f'item {i} is {item.value}')

[Icxolu]

I think you are looking for Py::borrow to get access to your container from within the iterator. Something like the following:

#[pymethods]
impl Container {
    ...

    fn __iter__(slf: Py<Self>) -> ContainerIterator {
        ContainerIterator {
            index: 0,
            container: slf,
        }
    }
}

#[pyclass]
struct ContainerIterator {
    index: usize,
    container: Py<Container>,
}

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

    fn __next__(&mut self, py: Python<'_>) -> PyResult<Option<Item>> {
        let container = self.container.borrow(py); // borrow to get access to the container
        if self.index < container.items.len() {
            let item = container.items[self.index].clone(); // needs clone, because we need to create new instances of `Item`
            self.index += 1;
            Ok(Some(item))
        } else {
            Ok(None)
        }
    }
}

One thing to note with this implementation is (in contrast to the example here), that modifications of the container, after creating the iterator, will be reflected by the iterator. This may or may not be desired. We still need to clone the items, because we need to store them in the Python instances that we are creating here. An alternative approach would be to hold Vec<Py<Item>> in your container, then your iterator can just Py::clone_ref each item, to increment the reference count of the existing instance instead of creating a new one. The trade-off is that Rust doesn't have direct access to the items anymore, so you always have to borrow them from Python before you can do something with them.

[aeshirey]

Thank you! self.container.borrow(py) does the trick. I was hoping to not have to clone my objects (and it looks like wrapping them in Arc isn't possible), but if that's the trade-off, then I can deal with it.