changeset 52878:c917656a259d

rust-pyo3-dirstate: CopyMap iterators This requires yet another `Sync` marker in `hg-core`. Of course we are leveraging the new `py_shared_iterator!` for this.
author Georges Racinet <georges.racinet@cloudcrane.io>
date Tue, 04 Feb 2025 13:36:42 +0100
parents 6b38ff460f2a
children aa2cfeed65c9
files rust/hg-core/src/dirstate.rs rust/hg-pyo3/src/dirstate.rs rust/hg-pyo3/src/dirstate/copy_map.rs rust/hg-pyo3/src/dirstate/dirstate_map.rs
diffstat 4 files changed, 65 insertions(+), 4 deletions(-) [+]
line wrap: on
line diff
--- a/rust/hg-core/src/dirstate.rs	Tue Feb 04 10:42:26 2025 +0100
+++ b/rust/hg-core/src/dirstate.rs	Tue Feb 04 13:36:42 2025 +0100
@@ -53,6 +53,7 @@
 pub type CopyMapIter<'a> = Box<
     dyn Iterator<Item = Result<(&'a HgPath, &'a HgPath), DirstateV2ParseError>>
         + Send
+        + Sync
         + 'a,
 >;
 
--- a/rust/hg-pyo3/src/dirstate.rs	Tue Feb 04 10:42:26 2025 +0100
+++ b/rust/hg-pyo3/src/dirstate.rs	Tue Feb 04 13:36:42 2025 +0100
@@ -20,7 +20,7 @@
     DirstateMapKeysIterator,
 };
 mod copy_map;
-use copy_map::CopyMap;
+use copy_map::{CopyMap, CopyMapItemsIterator, CopyMapKeysIterator};
 
 pub fn init_module<'py>(
     py: Python<'py>,
@@ -35,5 +35,7 @@
     m.add_class::<DirstateMapKeysIterator>()?;
     m.add_class::<DirstateMapItemsIterator>()?;
     m.add_class::<CopyMap>()?;
+    m.add_class::<CopyMapKeysIterator>()?;
+    m.add_class::<CopyMapItemsIterator>()?;
     Ok(m)
 }
--- a/rust/hg-pyo3/src/dirstate/copy_map.rs	Tue Feb 04 10:42:26 2025 +0100
+++ b/rust/hg-pyo3/src/dirstate/copy_map.rs	Tue Feb 04 13:36:42 2025 +0100
@@ -10,11 +10,17 @@
 
 use pyo3::exceptions::PyKeyError;
 use pyo3::prelude::*;
-use pyo3::types::{PyBytes, PyDict};
+use pyo3::types::{PyBytes, PyDict, PyTuple};
+use pyo3_sharedref::py_shared_iterator;
 
 use std::sync::{RwLockReadGuard, RwLockWriteGuard};
 
-use hg::{dirstate::owning::OwningDirstateMap, utils::hg_path::HgPath};
+use hg::{
+    dirstate::{
+        on_disk::DirstateV2ParseError, owning::OwningDirstateMap, CopyMapIter,
+    },
+    utils::hg_path::HgPath,
+};
 
 use super::dirstate_map::DirstateMap;
 use crate::{
@@ -115,6 +121,18 @@
         })
     }
 
+    fn __iter__(&self, py: Python) -> PyResult<CopyMapKeysIterator> {
+        self.keys(py)
+    }
+
+    fn keys(&self, py: Python) -> PyResult<CopyMapKeysIterator> {
+        CopyMapKeysIterator::new(self.dirstate_map.bind(py))
+    }
+
+    fn items(&self, py: Python) -> PyResult<CopyMapItemsIterator> {
+        CopyMapItemsIterator::new(self.dirstate_map.bind(py))
+    }
+
     fn __setitem__(
         &self,
         py: Python,
@@ -146,7 +164,47 @@
     }
 }
 
+py_shared_iterator!(
+    CopyMapKeysIterator,
+    PyBytes,
+    DirstateMap,
+    inner,
+    CopyMapIter<'static>,
+    |dsm| dsm.copy_map_iter(),
+    CopyMap::keys_next_result
+);
+
+py_shared_iterator!(
+    CopyMapItemsIterator,
+    PyTuple,
+    DirstateMap,
+    inner,
+    CopyMapIter<'static>,
+    |dsm| dsm.copy_map_iter(),
+    CopyMap::items_next_result
+);
+
 impl CopyMap {
+    fn keys_next_result(
+        py: Python,
+        res: Result<(&HgPath, &HgPath), 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, &HgPath), DirstateV2ParseError>,
+    ) -> PyResult<Option<Py<PyTuple>>> {
+        let (key, value) = res.map_err(dirstate_v2_error)?;
+        Ok(Some(
+            (PyHgPathRef(key), PyHgPathRef(value))
+                .into_pyobject(py)?
+                .unbind(),
+        ))
+    }
+
     fn with_dirstate_map_read<T>(
         &self,
         py: Python,
--- a/rust/hg-pyo3/src/dirstate/dirstate_map.rs	Tue Feb 04 10:42:26 2025 +0100
+++ b/rust/hg-pyo3/src/dirstate/dirstate_map.rs	Tue Feb 04 13:36:42 2025 +0100
@@ -51,7 +51,7 @@
 
 #[pyclass(mapping)]
 pub struct DirstateMap {
-    inner: PyShareable<OwningDirstateMap>,
+    pub(super) inner: PyShareable<OwningDirstateMap>,
 }
 
 #[pymethods]