--- a/rust/hg-pyo3/src/dirstate/dirstate_map.rs	Thu Feb 06 11:18:28 2025 +0100
+++ b/rust/hg-pyo3/src/dirstate/dirstate_map.rs	Wed Jan 29 18:26:10 2025 +0100
@@ -8,20 +8,34 @@
 //! Bindings for the `hg::dirstate::dirstate_map` file provided by the
 //! `hg-core` package.
 
+use pyo3::exceptions::PyKeyError;
 use pyo3::prelude::*;
 use pyo3::types::{PyBytes, PyBytesMethods, PyTuple};
 use pyo3_sharedref::PyShareable;
 
-use hg::dirstate::{
-    dirstate_map::DirstateIdentity as CoreDirstateIdentity,
-    owning::OwningDirstateMap,
+use std::sync::RwLockReadGuard;
+
+use hg::{
+    dirstate::{
+        dirstate_map::{
+            DirstateIdentity as CoreDirstateIdentity, DirstateMapWriteMode,
+        },
+        owning::OwningDirstateMap,
+    },
+    utils::hg_path::HgPath,
+    DirstateParents,
 };
 
-use crate::{exceptions::dirstate_error, node::PyNode, utils::PyBytesDeref};
+use super::item::DirstateItem;
+use crate::{
+    exceptions::{dirstate_error, dirstate_v2_error, map_try_lock_error},
+    node::{node_from_py_bytes, PyNode},
+    path::PyHgPathRef,
+    utils::PyBytesDeref,
+};
 
-#[pyclass]
+#[pyclass(mapping)]
 pub struct DirstateMap {
-    #[allow(dead_code)]
     inner: PyShareable<OwningDirstateMap>,
 }
 
@@ -79,6 +93,147 @@
             inner: OwningDirstateMap::new_empty(vec![], None).into(),
         })
     }
+
+    #[pyo3(signature = (key, default=None))]
+    fn get(
+        slf: &Bound<'_, Self>,
+        key: &Bound<'_, PyBytes>,
+        default: Option<PyObject>,
+    ) -> PyResult<Option<PyObject>> {
+        let path = HgPath::new(key.as_bytes());
+
+        Self::with_inner_read(slf, |_self_ref, inner| {
+            match inner.get(path).map_err(dirstate_v2_error)? {
+                Some(entry) => Ok(Some(
+                    DirstateItem::new_as_py(slf.py(), entry)?.into_any(),
+                )),
+                None => Ok(default),
+            }
+        })
+    }
+
+    /// Returns suitable data for writing on disk in v1 format
+    ///
+    /// Despite the name, this is not a mutation of the object.
+    fn write_v1(
+        slf: &Bound<'_, Self>,
+        py: Python,
+        p1: &Bound<'_, PyBytes>,
+        p2: &Bound<'_, PyBytes>,
+    ) -> PyResult<Py<PyBytes>> {
+        Self::with_inner_read(slf, |_self_ref, inner| {
+            let parents = DirstateParents {
+                p1: node_from_py_bytes(p1)?,
+                p2: node_from_py_bytes(p2)?,
+            };
+            let packed = inner.pack_v1(parents).map_err(dirstate_error)?;
+            // TODO optim, see `write_v2()`
+            Ok(PyBytes::new(py, &packed).unbind())
+        })
+    }
+
+    /// Returns suitable new data for writing on disk in v2 format
+    ///
+    /// Despite the name, this is not a mutation of the object.
+    ///
+    /// The new data together with whether that data should be appended to
+    /// the existing data file whose content is at `self.on_disk` (True),
+    /// instead of written to a new data file (False).
+    fn write_v2(
+        slf: &Bound<'_, Self>,
+        py: Python,
+        write_mode: usize,
+    ) -> PyResult<Py<PyTuple>> {
+        Self::with_inner_read(slf, |_self_ref, inner| {
+            let rust_write_mode = match write_mode {
+                0 => DirstateMapWriteMode::Auto,
+                1 => DirstateMapWriteMode::ForceNewDataFile,
+                2 => DirstateMapWriteMode::ForceAppend,
+                _ => DirstateMapWriteMode::Auto, // XXX should we error out?
+            };
+            let (packed, tree_metadata, append, _old_data_size) =
+                inner.pack_v2(rust_write_mode).map_err(dirstate_error)?;
+            // TODO optim. In theory we should be able to avoid these copies,
+            // since we have full ownership of `packed` and `tree_metadata`.
+            // But the copy is done by CPython itself, in
+            // `PyBytes_FromStringAndSize()`. Perhaps something better can
+            // be done with `PyBytes_FromObject` (buffer protocol).
+            let packed = PyBytes::new(py, &packed).unbind();
+            let tree_metadata =
+                PyBytes::new(py, tree_metadata.as_bytes()).unbind();
+            Ok((packed, tree_metadata, append).into_pyobject(py)?.into())
+        })
+    }
+
+    fn __len__(slf: &Bound<'_, Self>) -> PyResult<usize> {
+        Self::with_inner_read(slf, |_self_ref, inner| Ok(inner.len()))
+    }
+
+    fn __contains__(
+        slf: &Bound<'_, Self>,
+        key: &Bound<'_, PyBytes>,
+    ) -> PyResult<bool> {
+        Self::with_inner_read(slf, |_self_ref, inner| {
+            inner
+                .contains_key(HgPath::new(key.as_bytes()))
+                .map_err(dirstate_v2_error)
+        })
+    }
+
+    fn __getitem__(
+        slf: &Bound<'_, Self>,
+        key: &Bound<'_, PyBytes>,
+    ) -> PyResult<Py<DirstateItem>> {
+        let key_bytes = key.as_bytes();
+        let path = HgPath::new(key_bytes);
+        Self::with_inner_read(slf, |_self_ref, inner| {
+            match inner.get(path).map_err(dirstate_v2_error)? {
+                Some(entry) => DirstateItem::new_as_py(slf.py(), entry),
+                None => Err(PyKeyError::new_err(
+                    String::from_utf8_lossy(key_bytes).to_string(),
+                )),
+            }
+        })
+    }
+
+    fn debug_iter(
+        slf: &Bound<'_, Self>,
+        py: Python,
+        all: bool,
+    ) -> PyResult<PyObject> {
+        Self::with_inner_read(slf, |_self_ref, inner| {
+            // the iterator returned by `debug_iter()` does not
+            // implement ExactSizeIterator, which is needed by
+            // `PyList::new()`, so we need to collect. Probably not a
+            // performance issue, as this is a debug method.
+            let as_vec: PyResult<Vec<_>> = inner
+                .debug_iter(all)
+                .map(|item| {
+                    let (path, (state, mode, size, mtime)) =
+                        item.map_err(dirstate_v2_error)?;
+                    Ok((PyHgPathRef(path), state, mode, size, mtime))
+                })
+                .collect();
+            Ok(as_vec?.into_pyobject(py)?.unbind())
+        })
+    }
+}
+
+impl DirstateMap {
+    fn with_inner_read<'py, T>(
+        slf: &Bound<'py, Self>,
+        f: impl FnOnce(
+            &PyRef<'py, Self>,
+            RwLockReadGuard<OwningDirstateMap>,
+        ) -> PyResult<T>,
+    ) -> PyResult<T> {
+        let self_ref = slf.borrow();
+        // Safety: the owner is the right one. We will anyway
+        // not actually `share` it.
+        let shareable_ref = unsafe { self_ref.inner.borrow_with_owner(slf) };
+        let guard = shareable_ref.try_read().map_err(map_try_lock_error)?;
+        f(&self_ref, guard)
+    }
 }
 
 #[pyclass]

--- a/rust/hg-pyo3/src/exceptions.rs	Thu Feb 06 11:18:28 2025 +0100
+++ b/rust/hg-pyo3/src/exceptions.rs	Wed Jan 29 18:26:10 2025 +0100
@@ -92,7 +92,6 @@
     PyOSError::new_err(format!("Dirstate error: {:?}", err))
 }
 
-#[allow(dead_code)]
 pub fn dirstate_v2_error(_err: DirstateV2ParseError) -> PyErr {
     PyValueError::new_err("corrupted dirstate-v2")
 }

--- a/rust/hg-pyo3/src/path.rs	Thu Feb 06 11:18:28 2025 +0100
+++ b/rust/hg-pyo3/src/path.rs	Wed Jan 29 18:26:10 2025 +0100
@@ -15,7 +15,6 @@
 
 use hg::utils::hg_path::{HgPath, HgPathBuf};
 
-#[allow(dead_code)]
 #[derive(Eq, Ord, PartialEq, PartialOrd, Hash, derive_more::From)]
 pub struct PyHgPathRef<'a>(pub &'a HgPath);