pyo3
pyo3 copied to clipboard
PyO3
Iterator examples without `Clone`
As originally posted on Gitter by @microgravitas:
Hi y'all! Can anyone point me to an example of how to implement an Python iterator which does not copy the whole underlying datastructure? All examples I could find used .clone(), something I'd like to avoid---if possible.
I've actually noticed this too about our iterator examples - it would be nice if we could find a way to make the .clone() unnecessary!
There is obviously no way when using vec::IntoIter, but that is also not really necessary.
The iterator type should just need a reference to the Py object, and an index.
Pretty much. I was thinking similar to the Python standard library types we might also want a guard to check that the base object is not modified during iteration. I think that'd have to be built into the base object, maybe as an additional field.
To be clear, is this issue about:
- using a Python Iterator from Rust, or
- using a Rust Iterator from Python?
If this issue isn't about using Rust Iterator from Python, what would be the best way to do this? In other words, can anyone answer this question?
From what I can tell, the problem is that iterator types returning references require a lifetime. And that lifetime propagates to any type that holds it. In PyO3, it's not possible for a struct to use a lifetime, as noted in this issue.
This is exactly about using a "Rust iterator from Python". The example I'm thinking of is this one in the guide, which uses .clone().into_iter(), which is essentially just as bad as the stackoverflow link you post.
Yes, I've seen that one from the guide - thanks for the confirmation.
Is the problem that such a Rust Iterator to be used from Python (without clone) isn't possible, or that it's just not documented currently?
If it is possible, can you point me in the general direction of how to do it?
It should be possible! From the discussion between myself and birkenfeld above, I think this might be the pair of types that you'd want:
#[pyclass]
struct Container {
vec: Vec<usize>,
// if you want to reproduce Python's "X modified during iteration" behaviour, you'll probably need
// a flag like this to:
// - set it false whenever iteration starts
// - set it true every time the vec is modified
// - check it every time Iter.__next__ is called
has_been_modified: bool,
}
#[pyclass]
struct Iter {
src: Py<Container>,
index: usize
}
I've left the implementation of the protocols out - it should be similar to the example I linked; hopefully you can figure it out from the types I've written. Feel free to ask further questions if what I've just written isn't clear enough!
If you manage to get this working, would you be willing to provide a PR to update the linked example? 🙂
@davidhewitt That looks like a fine solution for indexable data structures like vectors, but is there any chance of doing this for arbitrary rust iterators? I bumped into exactly the issue @rapodaca mentioned: it seems impossible to avoid adding a lifetime to anything that holds a reference to an iterator.
For the more general case - the answer is that it's really really hard. PyO3 might be able to one day be able to do this, but it'll require a lot of careful R&D first.
See #1088 and #1089 for some thoughts I've had on this topic this morning.
My use cases appear to be the tough ones: wrap an existing iterator; or build an iterator to a non-indexable data structure such as HashSet.
#1089 is interesting. If I'm reading it right the idea is to build a Python iterator from Rust.
What about PyIterator? I'm not sure if it's still part of the project, but the line: " this class includes a Python<'p> token so that PyIterator can implement the rust Iterator trait" sounds encouraging.
I also found this which contains an example that isn't present anymore.
The documentation here is quite confusing. It's hard to know if what I'm looking at is out of date or current. Or even if it goes in the right direction.
PyIterator is for the reverse direction - it's equivalent to calling iter(foo) on a Python object so that you can then iterate its contents.
And it's still part of the current version of pyo3 😄 https://docs.rs/pyo3/0.11.1/pyo3/types/struct.PyIterator.html
I'm sorry I don't follow all the discussion, but please be aware that all elements must be already Python object if a container is non-copy iterator. E.g., :
#[pyclass]
struct User {
name: String,
id: usize,
}
#[pyclass]
struct UserIterator {
contents: Vec<Py<User>>,
index: usize,
}
#[pyproto]
impl<'p> PyIterProtocol for UserIterator {
fn __iter__(slf: PyRef<'p, Self>) -> PyRef<'p, Self> {
slf
}
fn __next__(mut slf: PyRefMut<'p, Self>) -> Option<Py<User>> {
let index = slf.index;
slf.index += 1;
slf.contents.get(index).map(|user| user.clone_ref(slf.py()))
}
}
I started to build a proof of concept for a "borrowed iterator" using ouroboros. I have some working code for a specific iterable struct and I'm trying to generalize it a bit. Here is a description of the mechanism I want to implement.
I defined a custom "IteratorIntoPy" to expose a "next_py" function returning an Option<PyObject>. A struct containing a boxed instance of this trait is exposed as a Python class and provides __iter__ and __next__ methods. Maybe this could reuse an existing trait from PyO3 instead?
Now the core of the implementation is a generic struct (called BorrowedIterator) using ouroboros to hold:
- a copy of the parent Python object
- a borrowed PyRef of this parent object
- an iterator borrowing further from this PyRef
As this struct implements the previous "IteratorIntoPy" trait, it can be wrapped into our generic boxed Python object.
Now for a rust type T exposed as a Python class, we can implement __iter__ creating a boxed BorrowedIterator providing the Py<T> object and a closure to extract the iterator.
Preliminary code works on my first use case, and as it holds an open PyRef, it blocks all concurrent calls to functions of "&mut self" (which I consider good behaviour).
As a side effect, these calls are blocked until the iterator is fully released, which may or may not happen in a predictable way. This could be improved in two ways:
- only open the PyRef on the first call to
__next__ - close the PyRef when the iterator is finished (or on demand with a close method).
This would not eliminate all errors in legitimate situations (e.g. call a mutable method before consuming the iterator but after it's last actual use) but it would make it good enough for most use cases.
Another limitation is that we need to returned owned Python objects, thus clone the returned items. In my use case it's fine (they are small and short-lived objects), but you may want to return an object that also borrows from the parent, thus maintains an open PyRef. This should be feasible but would trigger more competition and race conditions.
My current implementation uses a bit of unsafe code to extend the lifetime of the open PyRef to 'static. It's obviously not static, but ouroboros should handle this properly. I'm not really thrilled by this.
The wrapper iterator function returns PyObject instead of <T: IntoPy>, which requires to grab the GIL when __next__ is called and forward it to the next_py method. It's not ideal but it should not have a huge performance cost and the extra code is hidden behind the abstraction layer.
Does this make sense and could something like this make in into PyO3 itself? If so, I will be happy to provide some code and refactor it as needed.
@aurelien-naldi thank you for the detailed write-up and for spending time on this. If you are interested in working on this topic, please do! It's quite wide-ranging so I think it will take us a few tries to get the perfect design, hopefully we can make progress in stages.
Being able to safely handle borrowed data is something I'd really love to have PyO3 support. You might also want to read the discussions in #1089 and #2300 which have a bunch of related ideas in this space.
Some comments on your design in general:
Another limitation is that we need to returned owned Python objects, thus clone the returned items. In my use case it's fine (they are small and short-lived objects), but you may want to return an object that also borrows from the parent, thus maintains an open PyRef. This should be feasible but would trigger more competition and race conditions.
This is pretty much exactly what #1089 and #2300 are thinking about. On the competition specifically, I'm also wondering if in the future classes should be immutable, with individual fields protected by the GIL (as a kind of mutex), e.g.:
#[pyclass(immutable)] // immutable keyword may eventually be the default
struct ImmutableClass {
mutable_member: GILCell<i32>, // This field is readwrite, protected by the GIL
mutex_member: Mutex<Vec<u8>>, // This field is readwrite, protected by a mutex
}
Being able to lock individual fields might help avoid competition (by not having any access lock the whole object).
My current implementation uses a bit of unsafe code to extend the lifetime of the open PyRef to 'static. It's obviously not static, but ouroboros should handle this properly. I'm not really thrilled by this.
Yes, I think what is more "idiomatic" unsafe Rust is to break the PyRef down into raw pointers. The 'static lifetime is potentially very dangerous if the abstraction was accidentally leaky.
The wrapper iterator function returns PyObject instead of <T: IntoPy>, which requires to grab the GIL when next is called and forward it to the next_py method.
Until I've seen the code can't be sure, I'd hope that we can avoid the need to grab the GIL and just pass it in.
Thanks for the feedback!
I already had a look at #1089 and my current version is fairly close to the refmap proposed there, but uses Options instead of ManuallyDrop to store borrowed objects and adds an automatic release when the iterator is consumed.
I have a few open questions
Garbage collection
The borrowed object relies on a copy of an existing Python object, I'm a bit naive about how PyO3 handles this, but as I understand this relies on __traverse__ and __clear__, right?
Can we rely on the fact that we implement them and always export this borrowed object to Python so that these methods are properly called? We can also hide the copied object behind an Option and adapt the drop implementation to grabs the GIL if the object is still there. Am I missing other corner cases?
In any case, I would rather avoid exposing the underlying struct for now and only provide a wrapper providing the right pyclass methods (e.g. the iterator wrapper currently available).
Ideally we could later provide a more generic API to wrap a borrowed object in a Python object ensuring all the memory handling and delegating other function calls to an helper object. This may still require exposing some dangerous data structures, but maybe they can be marked as unsafe and hidden behind safe-looking macros?
Handling competition
I also think that we can better handle the case where a mutable function is called while the borrowed object still exists and only block the next access in the borrowed object. But hiding the borrowed reference behind a mutex seems too restrictive as it means that the owner object need to keep track of all ongoing borrows. At the PyO3 level, I think we can put in place a generic mechanism to prevent concurrent modification at the object level and let end-users handle more subtle cases themselves. Note that I'm drafting that at the PyCell level, but it's just a guess, please let me know how wrong it is!
struct PyCell {
// ... existing fields ...
grab: Option<Arc<AtomicBool>>
}
impl PyCell {
/// Declare that someone wants to grab this object
/// Return a Arc<AtomicBool> that remains true until the next call to a mutable function
/// Callers know that while this grab is locked and true the underlying object has not been changed.
/// TODO: is an AtomicBool enough or do we need a Mutex<bool> to properly ensure locking ?
fn grab(&mut self) -> Arc<AtomicBool>;
/// When a mutable PyRef is obtained, it checks if a grab exists.
/// In absence of grab, it can continue safely (it holds the pyref, so no call to grab() can happen)
/// In presence of grab:
/// * set it to false (to warn grab holders)
/// * forget the grab (to avoid wasting more time on it and wrongly setting it back to true)
/// Something similar could be done for grab_mut, but it's not needed for regular iterators
}
With something like this in place, I think that we can replace the locked PyRef with an unsafe pointer that can only be accessed as long as the grab is valid. Unless there is a risk that the pointer becomes invalid even without any mutable access?
I understand this relies on traverse and clear, right?
These are only relevant if you have cycles in your struct that the garbage collector must break. If you don't have cycles, you don't need them.
Handling competition
I'd prefer to see some concrete test cases written out which demonstrates a bunch of these "borrowed objects / borrowed iterators" and how they interact with the existing cell machinery before we try to hash out the implementation details too much. I'd quite like to define the target in terms of the API / examples that we would document for users.
I also think that we don't need to do anything too complex, because the "fine-grained" control for users is the just-merged #[pyclass(immutable)] where the objects themselves don't have any locking at all. Users can then wrap individual fields in Mutex (or GilCell which I'm about to propose) to control how constituent parts lock.
The immutable objects with separate locks for individual fields sounds great, but I'm not convinced that it helps much with the iterator use case (which is a separate object accessing borrowed data).
I started to write some example code (outside of Pyo3) to illustrate how I think we could handle a borrowed pointer that remains valid as long as the owner object is unmodified without locking it, but it's full of subtle moving pieces, thus I agree that without carefully looking at the code we can't really discuss. When I have something working, should I post it here or in #1089 ?
Here is fine, as that's where the discussion is. Even better might be to open as a PR so that we can comment inline.
Here is a first draft: https://gist.github.com/aurelien-naldi/905162889f2a3d90c09fc72f5b74cb82
It has a lot of boilerplate as it duplicates the core PyO3 structs (PyCell and associated references) but is self-contained.
Depending on your feedback, I can ditch this, turn it into its own playground repository to refine the idea or upgrade it into a fork of PyO3 and send a PR (improving error handling and interaction with the GIL).
If we are happy with the core borrowing mechanism, it should be "easy" to wrap it into generic iterators (exposed to Python directly) if the items are either small "Copy" structs (my main use case) or already Python objects.
Thanks @auerlien-naldi! I'd like to take some time to play around with this for a bit and collect my thoughts. Will try to get back with feedback within a week or so.
Thanks @auerlien-naldi! I'd like to take some time to play around with this and collect my thoughts. Will try to get back with feedback within a week or so.
With apologies I haven't forgotten but also been stretched very thin so haven't had a chance to sit down and give this some quality time. It's still on my radar, although please be patient :)
I've had to implement this functionality (for mutable iterators over some Container), and this is what I've learned:
- The data stored in a Container is usually some Rust struct (and not a PyObject of any kind). For example, our Container could be a sort of Vec-like storage over a struct with a couple of u32 fields. Typically, if you store PyObjects, you would not implement the datastructure in Rust. When storing Rust objects, the BorrowingCell does not quite work.
- Further issue is that the iterator over underlying Rust datastructure will have some lifetime constraints which prevent it from being inserted into a pyclass.
- More fundamentally, the Container and Iterator pyclasses both want to have mutable access to the underlying data (container for direct access via things like setitem and iterator for normal slice access). This means that when Iterator is created, the ownership of data must be transferred to it to avoid contention/invalidation (or some sort of lockout put onto the Container). The easiest way to do that is to store
Option<Box<T>>in the Container, and literally take it when making Iterator. It helps if T is Pin, this way you can be sure it does not get moved around. If someone tries to mess with Container while iterating, exceptions can be raised. - When Iterator is destroyed (i.e. its Drop is called) we can put back the data into Container. The easiest way I've found to do that is to have something like
Arc<RwLock<Option<Box<T>>>as the stored type in both the Container and Iterator. Then the data field can be "returned" to container without having to hold a reference to Container itself as a Pyclass. Not sure if that is such a brilliant idea though, as there is likely a way to reuse GIL as the RwLock in this case. - Ouroboros is an easy and straightforward (if a bit awkward) way to get around all the lifetime constraints related to storage of Rust iterators. Lifetime extension is also a viable option but I have not tried it quite yet. Either way Iterator needs to have a way to erase lifetimes from Rust iterator over underlying datastructure such that it can be stored in a pyclass. Semantically, storing the iterator and struct over which it is running in the same place makes sense, and this could probably become its own crate (with one template function that simply wraps up whatever container and whatever iterator over that same container into a self-referential structure without lifetime parameters).
- The Iterator pyclass in turn hands out mutable references to pieces of underlying rust data structure. At this point using
Arc<RwLock <...>>sort of logic for each accessed element becomes too cumbersome, and something else should be done to preserve performance and sanity. - Solution I've used is to make a Cursor class that holds a raw (unsafe) pointer into underlying Rust data structure, and has get/set methods to modify it via that pointer. This way Python code can mutate the contents of Rust structs that are still sitting inside the Rust container, no clones made. Naturally, if someone was to drop the whole Container at this point, the world would explode. To avoid this 100%, Cursor needs to hold a reference to the Iterator (keeping it from getting dropped), which in turn prevents the data block from disappearing.
- To avoid recreating the Cursor on every call to next (and thus having the associated pressure on memory allocation), we can instead update the raw pointer inside already existing Cursor instance, and return it again and again. To make sure Python did not leak references to the Cursor on a previous iteration, we check that Cursor's refcount is exactly 1 (with the Iterator holding the reference) every time next is called.
- In some ideal world a procmacro can be made to implement Cursor over any sane Rust struct automatically.
All of these ideas do look somewhat unrefined to me, but I'm at a loss as to what can be done better. This further does not cover the case where multiple iterators over the same Container may make reads without any issues (and thus when those are made there is no need to move the data away from Container). However, it does cover the key needs:
- performant mutable iteration over Rust backing storage with in-place edits from Python
- reasonable level of memory safety
- container lock-out during iteration
I haven't seen a practical working example so I will post my poor solution here in case it helps others.
This is not generic, I am just trying to get something working for my very specific use case of iterating over BTreeMap.keys.
The main idea is that I have a pyclass classe PyBTreeMap that wraps a BTreeMap. PyBTreeKeyIterator is created when you call the keys methods. PyBTreeKeyIterator has 2 fields:
-
owner: holds a reference to thePyObjectwhich created it and holds the data, you literally don't use it but having it makes sure the underlying data is not dropped. -
iter: the inner iterator, in this case of typebtree_map::Keys.
#[pyclass]
pub struct PyBTreeKeyIterator {
#[pyo3(get)]
pub owner: PyObject,
iter: btree_map::Keys<'static, PyObjectWrapper, PyObjectWrapper>,
}
The most difficult thing is actually constructing this. You can do so by passing a &PyRef<'_, PyBTreeMap> (the one you get from the slf trick) to the new method. From there you create the owner: PyObject using into_py(), and the iter: btree_map::Keys by calling the keys(). The main rick/cheat is to pass iter into PyBTreeKeyIterator using unsafe + std::mem::transmute to force the lifetime as static.
impl PyBTreeKeyIterator {
pub fn new(py_btree: &PyRef<'_, PyBTreeMap>) -> Self {
return Python::with_gil(|py| {
let owner = py_btree.into_py(py);
let iter = py_btree.tree.keys();
return PyBTreeKeyIterator {
owner: owner.clone(),
iter: unsafe {
std::mem::transmute::<
btree_map::Keys<'_, PyObjectWrapper, PyObjectWrapper>,
btree_map::Keys<'static, PyObjectWrapper, PyObjectWrapper>,
>(iter)
},
};
});
}
}
The next method is pretty easy to implement once you have this. One last thing for season rustaseans, is there a chance of leaking memory with this method? So far I am just happy that it runs but I am not experienced enough to know if this is the best approach to unsafe code.
Here is the full script:
use std::collections::btree_map;
use crate::pybtree_map::PyBTreeMap;
use crate::pyobject_wrapper::PyObjectWrapper;
use pyo3::exceptions::PyStopIteration;
use pyo3::prelude::*;
#[pyclass]
pub struct PyBTreeKeyIterator {
#[pyo3(get)]
pub owner: PyObject,
iter: btree_map::Keys<'static, PyObjectWrapper, PyObjectWrapper>,
}
impl PyBTreeKeyIterator {
pub fn new(py_btree: &PyRef<'_, PyBTreeMap>) -> Self {
return Python::with_gil(|py| {
let owner = py_btree.into_py(py);
let iter = py_btree.tree.keys();
return PyBTreeKeyIterator {
owner: owner.clone(),
iter: unsafe {
std::mem::transmute::<
btree_map::Keys<'_, PyObjectWrapper, PyObjectWrapper>,
btree_map::Keys<'static, PyObjectWrapper, PyObjectWrapper>,
>(iter)
},
};
});
}
}
#[pymethods]
impl PyBTreeKeyIterator {
pub fn next(&mut self) -> PyResult<PyObject> {
match self.iter.next() {
Some(x) => Ok(x.obj.clone()),
None => Err(PyErr::new::<PyStopIteration, _>(())),
}
}
}
Small remark: There should be no need to call with_gil in the constructor, as a Python token is available via PyRef::py.
And a safety advice: You probably want to put the PyRef directly inside your iterator (also by transmuting the lifetime) to ensure that the PyCell holding the BTreeMap stays borrowed. As you wrote it, I think that Python code modify the underlying BTreeMap (by creating a new borrow of the PyCell) while you btree_map::Keys iterator is still alive which is undefined behaviour.
Thanks @adamreichold!
You probably want to put the PyRef directly inside your iterator (also by transmuting the lifetime) to ensure that the PyCell holding the BTreeMap stays borrowed
Q: is this better than having the PyObject or should I do both? A bit confused about ownership rules regarding PyObject vs PyRef. My understanding is that PyObject is managed by pythons reference counting so as I coded this, the PyBTreeMap holding the BTree would be kept alive by python.
There should be no need to call with_gil in the constructor, as a Python token is available via PyRef::py
Nice tip, removed a lot of with_gil calls using this.
Q: is this better than having the PyObject or should I do both? A bit confused about ownership rules regarding PyObject vs PyRef. My understanding is that PyObject is managed by pythons reference counting so as I coded this, the PyBTreeMap holding the BTree would be kept alive by python.
I think want both, the Py<PyObject> to keep at least single strong reference alive so the PyRef and the btree_map::Keys cannot become dangling and the PyRef itself so the PyCell storing the BTreeMap remains immutably borrowed.