rust-pyo3-revlog: _writinghandles
authorRapha?l Gom?s <rgomes@octobus.net>
Tue, 07 Jan 2025 17:57:58 +0100
changeset 52834 8ebe20a6fcb7
parent 52833 b7dd7af61488
child 52835 bc095c0db77c
rust-pyo3-revlog: _writinghandles
rust/hg-pyo3/src/revlog/mod.rs
--- a/rust/hg-pyo3/src/revlog/mod.rs	Sun Jan 05 23:20:14 2025 +0100
+++ b/rust/hg-pyo3/src/revlog/mod.rs	Tue Jan 07 17:57:58 2025 +0100
@@ -20,6 +20,7 @@
 use pyo3_sharedref::{PyShareable, SharedByPyObject};
 
 use std::collections::{HashMap, HashSet};
+use std::os::fd::AsRawFd;
 use std::sync::{
     atomic::{AtomicUsize, Ordering},
     RwLock, RwLockReadGuard, RwLockWriteGuard,
@@ -169,6 +170,37 @@
     }
 }
 
+// Only used from Python *tests*
+#[doc(hidden)]
+#[pyclass]
+pub struct PyFileHandle {
+    inner_file: std::os::fd::RawFd,
+}
+
+#[pymethods]
+impl PyFileHandle {
+    #[new]
+    fn new(handle: std::os::fd::RawFd) -> Self {
+        Self { inner_file: handle }
+    }
+
+    fn tell(&self, py: Python<'_>) -> PyResult<PyObject> {
+        let locals = PyDict::new(py);
+        locals.set_item("os", py.import("os")?)?;
+        locals.set_item("fd", self.inner_file)?;
+        let f = py.eval(c"os.fdopen(fd)", None, Some(&locals))?;
+
+        // Prevent Python from closing the file after garbage collecting.
+        // This is fine since Rust is still holding on to the actual File.
+        // (and also because it's only used in tests).
+        std::mem::forget(f.clone());
+
+        locals.set_item("f", f)?;
+        let res = py.eval(c"f.tell()", None, Some(&locals))?;
+        Ok(res.unbind())
+    }
+}
+
 #[pyclass]
 #[allow(dead_code)]
 struct InnerRevlog {
@@ -381,6 +413,40 @@
         })
     }
 
+    // This is only used in Python *tests*
+    #[getter]
+    #[doc(hidden)]
+    fn _writinghandles(
+        slf: &Bound<'_, Self>,
+        py: Python<'_>,
+    ) -> PyResult<PyObject> {
+        Self::with_core_read(slf, |_self_ref, irl| {
+            let handles = irl.python_writing_handles();
+            match handles.as_ref() {
+                None => Ok(py.None()),
+                Some(handles) => {
+                    let index_handle = PyFileHandle::new(
+                        handles.index_handle.file.as_raw_fd(),
+                    );
+                    let data_handle = handles
+                        .data_handle
+                        .as_ref()
+                        .map(|h| PyFileHandle::new(h.file.as_raw_fd()));
+                    Ok(PyTuple::new(
+                        py,
+                        &[
+                            index_handle.into_py_any(py)?,
+                            data_handle.into_py_any(py)?,
+                            py.None(), // Sidedata handle
+                        ],
+                    )?
+                    .unbind()
+                    .into())
+                }
+            }
+        })
+    }
+
     fn clear_cache(slf: &Bound<'_, Self>) -> PyResult<PyObject> {
         assert!(!Self::is_delaying(slf)?);
         let mut self_ref = slf.borrow_mut();