diff --git a/rust/hg-cpython/src/dirstate.rs b/rust/hg-cpython/src/dirstate.rs --- a/rust/hg-cpython/src/dirstate.rs +++ b/rust/hg-cpython/src/dirstate.rs @@ -11,21 +11,25 @@ //! From Python, this will be seen as `mercurial.rustext.dirstate` use cpython::{ - exc, PyBytes, PyDict, PyErr, PyInt, PyModule, PyObject, PyResult, - PySequence, PythonObject, PyTuple, Python, ToPyObject, + exc, ObjectProtocol, PyBytes, PyDict, PyErr, PyInt, PyModule, PyObject, + PyResult, PySequence, PyTuple, Python, PythonObject, ToPyObject, }; use hg::{ - pack_dirstate, parse_dirstate, CopyVecEntry, DirstateEntry, - DirstatePackError, DirstateParents, DirstateParseError, DirstateVec, + pack_dirstate, parse_dirstate, CopyVecEntry, DirsIterable, DirsMultiset, + DirstateEntry, DirstateMapError, DirstatePackError, DirstateParents, + DirstateParseError, DirstateVec, }; use std::collections::HashMap; use std::ffi::CStr; + #[cfg(feature = "python27")] extern crate python27_sys as python_sys; #[cfg(feature = "python3")] extern crate python3_sys as python_sys; + use self::python_sys::PyCapsule_Import; use libc::{c_char, c_int}; +use std::cell::RefCell; use std::mem::transmute; /// C code uses a custom `dirstate_tuple` type, checks in multiple instances @@ -102,20 +106,11 @@ } } -fn pack_dirstate_wrapper( +fn extract_dirstate_vec( py: Python, - dmap: PyDict, - copymap: PyDict, - pl: PyTuple, - now: PyInt, -) -> PyResult { - let p1 = pl.get_item(py, 0).extract::(py)?; - let p1: &[u8] = p1.data(py); - let p2 = pl.get_item(py, 1).extract::(py)?; - let p2: &[u8] = p2.data(py); - - let dirstate_vec: Result = dmap - .items(py) + dmap: &PyDict, +) -> Result { + dmap.items(py) .iter() .map(|(filename, stats)| { let stats = stats.extract::(py)?; @@ -136,7 +131,22 @@ }, )) }) - .collect(); + .collect() +} + +fn pack_dirstate_wrapper( + py: Python, + dmap: PyDict, + copymap: PyDict, + pl: PyTuple, + now: PyInt, +) -> PyResult { + let p1 = pl.get_item(py, 0).extract::(py)?; + let p1: &[u8] = p1.data(py); + let p2 = pl.get_item(py, 1).extract::(py)?; + let p2: &[u8] = p2.data(py); + + let dirstate_vec = extract_dirstate_vec(py, &dmap)?; let copies: Result, Vec>, PyErr> = copymap .items(py) @@ -150,7 +160,7 @@ .collect(); match pack_dirstate( - &dirstate_vec?, + &dirstate_vec, &copies?, DirstateParents { p1, p2 }, now.as_object().extract::(py)?, @@ -170,7 +180,10 @@ py, PyBytes::new(py, &filename[..]), decapsule_make_dirstate_tuple(py)?( - state as c_char, mode, size, mtime, + state as c_char, + mode, + size, + mtime, ), )?; } @@ -191,10 +204,103 @@ } } +py_class!(pub class Dirs |py| { + data dirs_map: RefCell; + + // `map` is either a `dict` or a flat iterator (usually a `set`, sometimes + // a `list`) + def __new__( + _cls, + map: PyObject, + skip: Option = None + ) -> PyResult { + let mut skip_state: Option = None; + if let Some(skip) = skip { + skip_state = Some(skip.extract::(py)?.data(py)[0] as i8); + } + let dirs_map; + + if let Ok(map) = map.cast_as::(py) { + let dirstate_vec = extract_dirstate_vec(py, &map)?; + dirs_map = DirsMultiset::new( + DirsIterable::Dirstate(dirstate_vec), + skip_state, + ) + } else { + let map: Result>, PyErr> = map + .iter(py)? + .map(|o| Ok(o?.extract::(py)?.data(py).to_owned())) + .collect(); + dirs_map = DirsMultiset::new( + DirsIterable::Manifest(map?), + skip_state, + ) + } + + Self::create_instance(py, RefCell::new(dirs_map)) + } + + def addpath(&self, path: PyObject) -> PyResult { + self.dirs_map(py).borrow_mut().add_path( + path.extract::(py)?.data(py), + ); + Ok(py.None()) + } + + def delpath(&self, path: PyObject) -> PyResult { + self.dirs_map(py).borrow_mut().delete_path( + path.extract::(py)?.data(py), + ) + .and(Ok(py.None())) + .or_else(|e| { + match e { + DirstateMapError::PathNotFound(_p) => { + Err(PyErr::new::( + py, + "expected a value, found none".to_string(), + )) + } + DirstateMapError::EmptyPath => { + Ok(py.None()) + } + } + }) + } + + // This is really inefficient on top of being ugly, but it's an easy way + // of having it work to continue working on the rest of the module + // hopefully bypassing Python entirely pretty soon. + def __iter__(&self) -> PyResult { + let dict = PyDict::new(py); + + for (key, value) in self.dirs_map(py).borrow().iter() { + dict.set_item( + py, + PyBytes::new(py, &key[..]), + value.to_py_object(py), + )?; + } + + let locals = PyDict::new(py); + locals.set_item(py, "obj", dict)?; + + py.eval("iter(obj)", None, Some(&locals)) + } + + def __contains__(&self, item: PyObject) -> PyResult { + Ok(self + .dirs_map(py) + .borrow() + .get(&item.extract::(py)?.data(py).to_owned()) + .is_some()) + } +}); + /// Create the module, with `__package__` given from parent pub fn init_module(py: Python, package: &str) -> PyResult { let dotted_name = &format!("{}.dirstate", package); let m = PyModule::new(py, dotted_name)?; + m.add(py, "__package__", package)?; m.add(py, "__doc__", "Dirstate - Rust implementation")?; m.add( @@ -219,6 +325,8 @@ ), )?; + m.add_class::(py)?; + let sys = PyModule::import(py, "sys")?; let sys_modules: PyDict = sys.get(py, "modules")?.extract(py)?; sys_modules.set_item(py, dotted_name, &m)?;