changeset 52873:c5773445d350

rust-pyo3-dirstate: DirstateMap iterators Pretty straightforward with `py_shared_iterator!`.
author Georges Racinet <georges.racinet@cloudcrane.io>
date Wed, 29 Jan 2025 19:13:46 +0100
parents 8f6d25439bdc
children 09eb477eec65
files rust/hg-core/src/dirstate.rs rust/hg-pyo3/src/dirstate.rs rust/hg-pyo3/src/dirstate/dirstate_map.rs
diffstat 3 files changed, 62 insertions(+), 3 deletions(-) [+]
line wrap: on
line diff
--- a/rust/hg-core/src/dirstate.rs	Tue Feb 04 11:55:14 2025 +0100
+++ b/rust/hg-core/src/dirstate.rs	Wed Jan 29 19:13:46 2025 +0100
@@ -46,6 +46,7 @@
     dyn Iterator<
             Item = Result<(&'a HgPath, DirstateEntry), DirstateV2ParseError>,
         > + Send
+        + Sync
         + 'a,
 >;
 
--- a/rust/hg-pyo3/src/dirstate.rs	Tue Feb 04 11:55:14 2025 +0100
+++ b/rust/hg-pyo3/src/dirstate.rs	Wed Jan 29 19:13:46 2025 +0100
@@ -15,7 +15,10 @@
 mod item;
 use item::DirstateItem;
 mod dirstate_map;
-use dirstate_map::{DirstateIdentity, DirstateMap};
+use dirstate_map::{
+    DirstateIdentity, DirstateMap, DirstateMapItemsIterator,
+    DirstateMapKeysIterator,
+};
 
 pub fn init_module<'py>(
     py: Python<'py>,
@@ -27,5 +30,7 @@
     m.add_class::<DirstateIdentity>()?;
     m.add_class::<DirstateItem>()?;
     m.add_class::<DirstateMap>()?;
+    m.add_class::<DirstateMapKeysIterator>()?;
+    m.add_class::<DirstateMapItemsIterator>()?;
     Ok(m)
 }
--- a/rust/hg-pyo3/src/dirstate/dirstate_map.rs	Tue Feb 04 11:55:14 2025 +0100
+++ b/rust/hg-pyo3/src/dirstate/dirstate_map.rs	Wed Jan 29 19:13:46 2025 +0100
@@ -11,7 +11,7 @@
 use pyo3::exceptions::PyKeyError;
 use pyo3::prelude::*;
 use pyo3::types::{PyBytes, PyBytesMethods, PyTuple};
-use pyo3_sharedref::PyShareable;
+use pyo3_sharedref::{py_shared_iterator, PyShareable};
 
 use std::sync::RwLockReadGuard;
 
@@ -20,7 +20,10 @@
         dirstate_map::{
             DirstateIdentity as CoreDirstateIdentity, DirstateMapWriteMode,
         },
+        entry::DirstateEntry,
+        on_disk::DirstateV2ParseError,
         owning::OwningDirstateMap,
+        StateMapIter,
     },
     utils::hg_path::HgPath,
     DirstateParents,
@@ -196,6 +199,18 @@
         })
     }
 
+    fn keys(slf: &Bound<'_, Self>) -> PyResult<DirstateMapKeysIterator> {
+        DirstateMapKeysIterator::new(slf)
+    }
+
+    fn items(slf: &Bound<'_, Self>) -> PyResult<DirstateMapItemsIterator> {
+        DirstateMapItemsIterator::new(slf)
+    }
+
+    fn __iter__(slf: &Bound<'_, Self>) -> PyResult<DirstateMapKeysIterator> {
+        Self::keys(slf)
+    }
+
     fn debug_iter(
         slf: &Bound<'_, Self>,
         py: Python,
@@ -214,13 +229,51 @@
                     Ok((PyHgPathRef(path), state, mode, size, mtime))
                 })
                 .collect();
+            // `IntoPyObject` on `Vec` and `&[T]` gives  `PyList` or `PyBytes`
             Ok(as_vec?.into_pyobject(py)?.unbind())
         })
     }
 }
 
+py_shared_iterator!(
+    DirstateMapKeysIterator,
+    PyBytes,
+    DirstateMap,
+    inner,
+    StateMapIter<'static>,
+    |dsm| dsm.iter(),
+    DirstateMap::keys_next_result
+);
+
+py_shared_iterator!(
+    DirstateMapItemsIterator,
+    PyTuple,
+    DirstateMap,
+    inner,
+    StateMapIter<'static>,
+    |dsm| dsm.iter(),
+    DirstateMap::items_next_result
+);
+
 impl DirstateMap {
-    fn with_inner_read<'py, T>(
+    fn keys_next_result(
+        py: Python,
+        res: Result<(&HgPath, DirstateEntry), DirstateV2ParseError>,
+    ) -> PyResult<Option<Py<PyBytes>>> {
+        let key = res.map_err(dirstate_v2_error)?.0;
+        Ok(Some(PyHgPathRef(key).into_pyobject(py)?.unbind()))
+    }
+
+    fn items_next_result(
+        py: Python,
+        res: Result<(&HgPath, DirstateEntry), DirstateV2ParseError>,
+    ) -> PyResult<Option<Py<PyTuple>>> {
+        let (key, entry) = res.map_err(dirstate_v2_error)?;
+        let py_entry = DirstateItem::new_as_py(py, entry)?;
+        Ok(Some((PyHgPathRef(key), py_entry).into_pyobject(py)?.into()))
+    }
+
+    pub fn with_inner_read<'py, T>(
         slf: &Bound<'py, Self>,
         f: impl FnOnce(
             &PyRef<'py, Self>,