changeset 52804:0ac956db7ea7

rust-pyo3-index: __getitem__
author Georges Racinet <georges.racinet@cloudcrane.io>
date Sun, 22 Dec 2024 21:37:29 +0100
parents 1b9907575768
children acae91fad6be
files rust/hg-pyo3/src/revlog/index.rs rust/hg-pyo3/src/revlog/mod.rs tests/test-rust-revlog.py
diffstat 3 files changed, 83 insertions(+), 6 deletions(-) [+]
line wrap: on
line diff
--- a/rust/hg-pyo3/src/revlog/index.rs	Tue Dec 24 15:08:22 2024 +0100
+++ b/rust/hg-pyo3/src/revlog/index.rs	Sun Dec 22 21:37:29 2024 +0100
@@ -39,3 +39,36 @@
         ..Default::default()
     })
 }
+
+pub fn revision_data_params_to_py_tuple(
+    py: Python<'_>,
+    params: RevisionDataParams,
+) -> PyResult<Bound<'_, PyTuple>> {
+    PyTuple::new(
+        py,
+        &[
+            params.data_offset.into_pyobject(py)?.into_any(),
+            params.data_compressed_length.into_pyobject(py)?.into_any(),
+            params
+                .data_uncompressed_length
+                .into_pyobject(py)?
+                .into_any(),
+            params.data_delta_base.into_pyobject(py)?.into_any(),
+            params.link_rev.into_pyobject(py)?.into_any(),
+            params.parent_rev_1.into_pyobject(py)?.into_any(),
+            params.parent_rev_2.into_pyobject(py)?.into_any(),
+            PyBytes::new(py, &params.node_id).into_any().into_any(),
+            params._sidedata_offset.into_pyobject(py)?.into_any(),
+            params
+                ._sidedata_compressed_length
+                .into_pyobject(py)?
+                .into_any(),
+            params.data_compression_mode.into_pyobject(py)?.into_any(),
+            params
+                ._sidedata_compression_mode
+                .into_pyobject(py)?
+                .into_any(),
+            params._rank.into_pyobject(py)?.into_any(),
+        ],
+    )
+}
--- a/rust/hg-pyo3/src/revlog/mod.rs	Tue Dec 24 15:08:22 2024 +0100
+++ b/rust/hg-pyo3/src/revlog/mod.rs	Sun Dec 22 21:37:29 2024 +0100
@@ -8,6 +8,8 @@
 // GNU General Public License version 2 or any later version.
 #![allow(non_snake_case)]
 use pyo3::buffer::PyBuffer;
+use pyo3::conversion::IntoPyObject;
+use pyo3::exceptions::PyIndexError;
 use pyo3::prelude::*;
 use pyo3::types::{PyBytes, PyBytesMethods, PyList, PyTuple};
 use pyo3_sharedref::PyShareable;
@@ -19,7 +21,7 @@
 
 use hg::{
     revlog::{
-        index::Index,
+        index::{Index, RevisionDataParams},
         inner_revlog::InnerRevlog as CoreInnerRevlog,
         nodemap::{NodeMap, NodeMapError, NodeTree as CoreNodeTree},
         options::RevlogOpenOptions,
@@ -27,7 +29,7 @@
     },
     utils::files::get_path_from_bytes,
     vfs::FnCacheVfs,
-    BaseRevision, Revision, UncheckedRevision,
+    BaseRevision, Revision, UncheckedRevision, NULL_REVISION,
 };
 
 use crate::{
@@ -44,7 +46,9 @@
 mod config;
 use config::*;
 mod index;
-use index::py_tuple_to_revision_data_params;
+use index::{
+    py_tuple_to_revision_data_params, revision_data_params_to_py_tuple,
+};
 
 #[pyclass]
 #[allow(dead_code)]
@@ -290,6 +294,41 @@
     fn _index___len__(slf: &Bound<'_, Self>) -> PyResult<usize> {
         Self::with_index_read(slf, |idx| Ok(idx.len()))
     }
+
+    fn _index___getitem__(
+        slf: &Bound<'_, Self>,
+        py: Python<'_>,
+        key: &Bound<'_, PyAny>,
+    ) -> PyResult<PyObject> {
+        Self::with_index_read(slf, |idx| {
+            match key.extract::<BaseRevision>() {
+                Ok(key_as_int) => {
+                    let entry_params = if key_as_int == NULL_REVISION.0 {
+                        RevisionDataParams::default()
+                    } else {
+                        let rev = UncheckedRevision(key_as_int);
+                        match idx.entry_as_params(rev) {
+                            Some(e) => e,
+                            None => {
+                                return Err(PyIndexError::new_err(
+                                    "revlog index out of range",
+                                ));
+                            }
+                        }
+                    };
+                    Ok(revision_data_params_to_py_tuple(py, entry_params)?
+                        .into_any()
+                        .unbind())
+                }
+                // Case when key is a binary Node ID (lame: we're re-unlocking)
+                _ => Self::_index_get_rev(slf, key.downcast::<PyBytes>()?)?
+                    .map_or_else(
+                        || Ok(py.None()),
+                        |py_rev| Ok(py_rev.into_pyobject(py)?.unbind().into()),
+                    ),
+            }
+        })
+    }
 }
 
 impl InnerRevlog {
--- a/tests/test-rust-revlog.py	Tue Dec 24 15:08:22 2024 +0100
+++ b/tests/test-rust-revlog.py	Sun Dec 22 21:37:29 2024 +0100
@@ -52,14 +52,19 @@
         idx = self.parserustindex()
         self.assertEqual(len(idx), 4)
 
+    def test_getitem(self):
+        idx = self.parserustindex()
+        as_tuple = (0, 82969, 484626, 0, 0, -1, -1, self.node0, 0, 0, 2, 2, -1)
+        self.assertEqual(idx[0], as_tuple)
+        self.assertEqual(idx[self.node0], 0)
+
     def test_index_append(self):
         idx = self.parserustindex(data=b'')
         self.assertEqual(len(idx), 0)
         self.assertIsNone(idx.get_rev(self.node0))
 
-        # this is the same first entry as in data provided by base test class
-        # (we do not have __getitem__ in the PyO3 version yet)
-        idx.append((0, 82969, 484626, 0, 0, -1, -1, self.node0, 0, 0, 2, 2, -1))
+        non_empty_index = self.parserustindex()
+        idx.append(non_empty_index[0])
         self.assertEqual(len(idx), 1)
         self.assertEqual(idx.get_rev(self.node0), 0)