rust-pyo3-dirstate: CopyMap without proxy methods
authorGeorges Racinet <georges.racinet@cloudcrane.io>
Tue, 04 Feb 2025 10:42:26 +0100
changeset 52865 6b38ff460f2a
parent 52864 d961e09d3d8c
child 52866 c917656a259d
rust-pyo3-dirstate: CopyMap without proxy methods It turns out that it is fairly easy to call the core `DirstateMap` methods from a `CopyMap` PyO3 object, hence we do not need to introduce proxy methods on the FFI `DirstateMap` as was done in `rust-cpython`. It is more straightforward, and is somewhat easier to make with PyO3 because we have all these intermediates between the Rust and Python layers (`Bound`, `Py<T>`). From the Python side, a difference with the `hg-cpython` version is that `CopyMap` would be instantiable also directly as `CopyMap(dirstate_map)`. Reproducing the `from_inner`, with refcount increase the responsibility in the `DirstateMap.copymap()` method would have been less natural (as there is no need with the separate `from_inner` in PyO3 constructors) and somehow we felt it was sounder to force the ref cloning from `CopyMap`.
rust/hg-pyo3/src/dirstate.rs
rust/hg-pyo3/src/dirstate/copy_map.rs
rust/hg-pyo3/src/dirstate/dirstate_map.rs
--- a/rust/hg-pyo3/src/dirstate.rs	Thu Jan 30 12:24:41 2025 +0100
+++ b/rust/hg-pyo3/src/dirstate.rs	Tue Feb 04 10:42:26 2025 +0100
@@ -19,6 +19,8 @@
     DirstateIdentity, DirstateMap, DirstateMapItemsIterator,
     DirstateMapKeysIterator,
 };
+mod copy_map;
+use copy_map::CopyMap;
 
 pub fn init_module<'py>(
     py: Python<'py>,
@@ -32,5 +34,6 @@
     m.add_class::<DirstateMap>()?;
     m.add_class::<DirstateMapKeysIterator>()?;
     m.add_class::<DirstateMapItemsIterator>()?;
+    m.add_class::<CopyMap>()?;
     Ok(m)
 }
--- /dev/null	Thu Jan 01 00:00:00 1970 +0000
+++ b/rust/hg-pyo3/src/dirstate/copy_map.rs	Tue Feb 04 10:42:26 2025 +0100
@@ -0,0 +1,167 @@
+// copy_map.rs
+//
+// Copyright 2019 Raphaël Gomès <rgomes@octobus.net>
+//           2025 Georges Racinet <georges.racinet@cloudcrane.io>
+//
+// This software may be used and distributed according to the terms of the
+// GNU General Public License version 2 or any later version.
+//! Bindings for `hg::dirstate::dirstate_map::CopyMap` provided by the
+//! `hg-core` package.
+
+use pyo3::exceptions::PyKeyError;
+use pyo3::prelude::*;
+use pyo3::types::{PyBytes, PyDict};
+
+use std::sync::{RwLockReadGuard, RwLockWriteGuard};
+
+use hg::{dirstate::owning::OwningDirstateMap, utils::hg_path::HgPath};
+
+use super::dirstate_map::DirstateMap;
+use crate::{
+    exceptions::dirstate_v2_error,
+    path::{PyHgPathBuf, PyHgPathRef},
+};
+
+#[pyclass(mapping)]
+pub struct CopyMap {
+    dirstate_map: Py<DirstateMap>,
+}
+
+#[pymethods]
+impl CopyMap {
+    #[new]
+    pub fn new(dsm: &Bound<'_, DirstateMap>) -> PyResult<Self> {
+        Ok(Self {
+            dirstate_map: dsm.clone().unbind(),
+        })
+    }
+
+    fn __getitem__(
+        &self,
+        py: Python,
+        key: &Bound<'_, PyBytes>,
+    ) -> PyResult<Py<PyBytes>> {
+        let key = key.as_bytes();
+        self.with_dirstate_map_read(py, |inner_dsm| {
+            inner_dsm
+                .copy_map_get(HgPath::new(key))
+                .map_err(dirstate_v2_error)?
+                .ok_or_else(|| {
+                    PyKeyError::new_err(
+                        String::from_utf8_lossy(key).to_string(),
+                    )
+                })
+                .and_then(|copy| {
+                    Ok(PyHgPathRef(copy).into_pyobject(py)?.unbind())
+                })
+        })
+    }
+
+    fn __len__(&self, py: Python) -> PyResult<usize> {
+        self.with_dirstate_map_read(py, |inner_dsm| {
+            Ok(inner_dsm.copy_map_len())
+        })
+    }
+
+    fn __contains__(
+        &self,
+        py: Python,
+        key: &Bound<'_, PyBytes>,
+    ) -> PyResult<bool> {
+        let key = key.as_bytes();
+        self.with_dirstate_map_read(py, |inner_dsm| {
+            inner_dsm
+                .copy_map_contains_key(HgPath::new(key))
+                .map_err(dirstate_v2_error)
+        })
+    }
+
+    #[pyo3(signature = (key, default=None))]
+    fn get(
+        &self,
+        py: Python,
+        key: &Bound<'_, PyBytes>,
+        default: Option<PyObject>,
+    ) -> PyResult<Option<PyObject>> {
+        let key = key.as_bytes();
+        self.with_dirstate_map_read(py, |inner_dsm| {
+            match inner_dsm
+                .copy_map_get(HgPath::new(key))
+                .map_err(dirstate_v2_error)?
+            {
+                Some(copy) => Ok(Some(
+                    PyHgPathRef(copy).into_pyobject(py)?.unbind().into(),
+                )),
+                None => Ok(default),
+            }
+        })
+    }
+
+    #[pyo3(signature = (key, default=None))]
+    fn pop(
+        &self,
+        py: Python,
+        key: &Bound<'_, PyBytes>,
+        default: Option<PyObject>,
+    ) -> PyResult<Option<PyObject>> {
+        let path = HgPath::new(key.as_bytes());
+        self.with_dirstate_map_write(py, |mut inner_dsm| {
+            match inner_dsm.copy_map_remove(path).map_err(dirstate_v2_error)? {
+                Some(copy) => Ok(Some(
+                    PyHgPathBuf(copy).into_pyobject(py)?.unbind().into(),
+                )),
+                None => Ok(default),
+            }
+        })
+    }
+
+    fn __setitem__(
+        &self,
+        py: Python,
+        key: &Bound<'_, PyBytes>,
+        value: &Bound<'_, PyBytes>,
+    ) -> PyResult<()> {
+        let key = HgPath::new(key.as_bytes());
+        let value = HgPath::new(value.as_bytes());
+        self.with_dirstate_map_write(py, |mut inner_dsm| {
+            inner_dsm
+                .copy_map_insert(key, value)
+                .map_err(dirstate_v2_error)
+        })?;
+        Ok(())
+    }
+
+    fn copy(&self, py: Python) -> PyResult<Py<PyDict>> {
+        let dict = PyDict::new(py);
+        // The `IntoPyDict` trait just does the same, but is not applicable
+        // here because it is meant to work on infallible iterators
+        self.with_dirstate_map_read(py, |inner_dsm| {
+            for item in inner_dsm.copy_map_iter() {
+                let (key, value) = item.map_err(dirstate_v2_error)?;
+                dict.set_item(PyHgPathRef(key), PyHgPathRef(value))?;
+            }
+            Ok(())
+        })?;
+        Ok(dict.unbind())
+    }
+}
+
+impl CopyMap {
+    fn with_dirstate_map_read<T>(
+        &self,
+        py: Python,
+        f: impl FnOnce(RwLockReadGuard<OwningDirstateMap>) -> PyResult<T>,
+    ) -> PyResult<T> {
+        let dsm = self.dirstate_map.bind(py);
+        DirstateMap::with_inner_read(dsm, |_dsm, inner| f(inner))
+    }
+
+    fn with_dirstate_map_write<T>(
+        &self,
+        py: Python,
+        f: impl FnOnce(RwLockWriteGuard<OwningDirstateMap>) -> PyResult<T>,
+    ) -> PyResult<T> {
+        let dsm = self.dirstate_map.bind(py);
+        DirstateMap::with_inner_write(dsm, |_dsm, inner| f(inner))
+    }
+}
--- a/rust/hg-pyo3/src/dirstate/dirstate_map.rs	Thu Jan 30 12:24:41 2025 +0100
+++ b/rust/hg-pyo3/src/dirstate/dirstate_map.rs	Tue Feb 04 10:42:26 2025 +0100
@@ -32,7 +32,7 @@
     DirstateParents,
 };
 
-use super::item::DirstateItem;
+use super::{copy_map::CopyMap, item::DirstateItem};
 use crate::{
     exceptions::{
         dirstate_error, dirstate_v2_error, map_try_lock_error,
@@ -383,6 +383,10 @@
         Self::keys(slf)
     }
 
+    fn copymap(slf: &Bound<'_, Self>) -> PyResult<Py<CopyMap>> {
+        CopyMap::new(slf).and_then(|cm| Py::new(slf.py(), cm))
+    }
+
     fn tracked_dirs(
         slf: &Bound<'_, Self>,
         py: Python,
@@ -478,7 +482,7 @@
         Ok(Some((PyHgPathRef(key), py_entry).into_pyobject(py)?.into()))
     }
 
-    pub fn with_inner_read<'py, T>(
+    pub(super) fn with_inner_read<'py, T>(
         slf: &Bound<'py, Self>,
         f: impl FnOnce(
             &PyRef<'py, Self>,
@@ -493,7 +497,7 @@
         f(&self_ref, guard)
     }
 
-    fn with_inner_write<'py, T>(
+    pub(super) fn with_inner_write<'py, T>(
         slf: &Bound<'py, Self>,
         f: impl FnOnce(
             &PyRef<'py, Self>,