changeset 52881:8c11ec902e73

rust-pyo3-dirstate: DirsMultisetKeysIterator
author Georges Racinet <georges.racinet@cloudcrane.io>
date Wed, 05 Feb 2025 11:20:39 +0100
parents 8e94b32de84b
children 38e16da74aea
files rust/hg-pyo3/src/dirstate.rs rust/hg-pyo3/src/dirstate/dirs_multiset.rs
diffstat 2 files changed, 26 insertions(+), 3 deletions(-) [+]
line wrap: on
line diff
--- a/rust/hg-pyo3/src/dirstate.rs	Wed Feb 05 11:11:05 2025 +0100
+++ b/rust/hg-pyo3/src/dirstate.rs	Wed Feb 05 11:20:39 2025 +0100
@@ -22,7 +22,7 @@
 mod copy_map;
 use copy_map::{CopyMap, CopyMapItemsIterator, CopyMapKeysIterator};
 mod dirs_multiset;
-use dirs_multiset::Dirs;
+use dirs_multiset::{Dirs, DirsMultisetKeysIterator};
 
 pub fn init_module<'py>(
     py: Python<'py>,
@@ -40,5 +40,6 @@
     m.add_class::<CopyMapKeysIterator>()?;
     m.add_class::<CopyMapItemsIterator>()?;
     m.add_class::<Dirs>()?;
+    m.add_class::<DirsMultisetKeysIterator>()?;
     Ok(m)
 }
--- a/rust/hg-pyo3/src/dirstate/dirs_multiset.rs	Wed Feb 05 11:11:05 2025 +0100
+++ b/rust/hg-pyo3/src/dirstate/dirs_multiset.rs	Wed Feb 05 11:20:39 2025 +0100
@@ -10,16 +10,17 @@
 use pyo3::exceptions::PyTypeError;
 use pyo3::prelude::*;
 use pyo3::types::{PyBytes, PyDict};
-use pyo3_sharedref::PyShareable;
+use pyo3_sharedref::{py_shared_iterator, PyShareable};
 
 use std::sync::{RwLockReadGuard, RwLockWriteGuard};
 
 use hg::{
-    dirstate::dirs_multiset::DirsMultiset,
+    dirstate::dirs_multiset::{DirsMultiset, DirsMultisetIter},
     utils::hg_path::{HgPath, HgPathBuf},
 };
 
 use crate::exceptions::{map_try_lock_error, to_string_value_error};
+use crate::path::PyHgPathRef;
 
 #[pyclass(mapping)]
 pub struct Dirs {
@@ -68,6 +69,10 @@
         })
     }
 
+    fn __iter__(slf: &Bound<'_, Self>) -> PyResult<DirsMultisetKeysIterator> {
+        DirsMultisetKeysIterator::new(slf)
+    }
+
     fn __contains__(
         slf: &Bound<'_, Self>,
         key: &Bound<'_, PyAny>,
@@ -82,7 +87,24 @@
     }
 }
 
+py_shared_iterator!(
+    DirsMultisetKeysIterator,
+    PyBytes,
+    Dirs,
+    inner,
+    DirsMultisetIter<'static>,
+    |ms| ms.iter(),
+    Dirs::keys_next_result
+);
+
 impl Dirs {
+    fn keys_next_result(
+        py: Python,
+        res: &HgPathBuf,
+    ) -> PyResult<Option<Py<PyBytes>>> {
+        Ok(Some(PyHgPathRef(res).into_pyobject(py)?.unbind()))
+    }
+
     pub(super) fn with_inner_read<T>(
         slf: &Bound<'_, Self>,
         f: impl FnOnce(RwLockReadGuard<DirsMultiset>) -> PyResult<T>,