Skip to content

Commit

Permalink
And now an iterator for HashTrieSet.
Browse files Browse the repository at this point in the history
This I'm even less sure is 'optimal' (well it's definitely not fully
but even for my level of optimal) -- but I don't see another way
to do 'give me an arbitrary element and then the remaining set'
other than creating an Iter.
  • Loading branch information
Julian committed Jan 10, 2024
1 parent e2f2842 commit 579e1ca
Showing 1 changed file with 22 additions and 8 deletions.
30 changes: 22 additions & 8 deletions src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -338,14 +338,10 @@ impl HashTrieSetPy {
self.symmetric_difference(other)
}

fn __iter__(slf: PyRef<'_, Self>) -> PyResult<Py<KeyIterator>> {
let iter = slf
.inner
.iter()
.map(|k| k.to_owned())
.collect::<Vec<_>>()
.into_iter();
Py::new(slf.py(), KeyIterator { inner: iter })
fn __iter__(slf: PyRef<'_, Self>) -> SetIterator {
SetIterator {
inner: slf.inner.clone(),
}
}

fn __len__(&self) -> usize {
Expand Down Expand Up @@ -491,6 +487,24 @@ impl HashTrieSetPy {
}
}

#[pyclass(module = "rpds")]
struct SetIterator {
inner: HashTrieSetSync<Key>,
}

#[pymethods]
impl SetIterator {
fn __iter__(slf: PyRef<'_, Self>) -> PyRef<'_, Self> {
slf
}

fn __next__(mut slf: PyRefMut<'_, Self>) -> Option<Key> {
let first = slf.inner.iter().next()?.to_owned();
slf.inner = slf.inner.remove(&first);
Some(first)
}
}

#[repr(transparent)]
#[pyclass(name = "List", module = "rpds", frozen, sequence)]
struct ListPy {
Expand Down

0 comments on commit 579e1ca

Please sign in to comment.