diff --git a/rust/treedirstate/src/dirstate.rs b/rust/treedirstate/src/dirstate.rs --- a/rust/treedirstate/src/dirstate.rs +++ b/rust/treedirstate/src/dirstate.rs @@ -186,6 +186,15 @@ self.tracked.get(self.store.store_view(), name) } + /// Visit all tracked files with a visitor. + pub fn visit_tracked(&mut self, visitor: &mut F) -> Result<()> + where + F: FnMut(&Vec, &T) -> Result<()>, + { + self.store.cache()?; + self.tracked.visit(self.store.store_view(), visitor) + } + /// Get the name and state of the first file in the tracked tree. pub fn get_first_tracked<'a>(&'a mut self) -> Result> { self.store.cache()?; @@ -206,6 +215,14 @@ self.removed.get(self.store.store_view(), name) } + /// Visit all removed files with a visitor. + pub fn visit_removed(&mut self, visitor: &mut F) -> Result<()> + where + F: FnMut(&Vec, &T) -> Result<()>, + { + self.removed.visit(self.store.store_view(), visitor) + } + /// Get the name and state of the first file in the removed tree. pub fn get_first_removed<'a>(&'a mut self) -> Result> { self.removed.get_first(self.store.store_view()) diff --git a/rust/treedirstate/src/errors.rs b/rust/treedirstate/src/errors.rs --- a/rust/treedirstate/src/errors.rs +++ b/rust/treedirstate/src/errors.rs @@ -20,6 +20,10 @@ CorruptTree { description("treedirstate is corrupt"), } + CallbackError(desc: String) { + description("callback error"), + display("callback error: {}", desc), + } } foreign_links { Io(::std::io::Error); diff --git a/rust/treedirstate/src/python.rs b/rust/treedirstate/src/python.rs --- a/rust/treedirstate/src/python.rs +++ b/rust/treedirstate/src/python.rs @@ -4,10 +4,12 @@ use cpython::*; use cpython::exc; use dirstate::Dirstate; +use errors::ErrorKind; use filestate::FileState; use std::cell::RefCell; use std::path::PathBuf; use store::BlockId; +use tree::KeyRef; py_module_initializer!( rusttreedirstate, @@ -19,6 +21,13 @@ } ); +fn callback_error(py: Python, mut e: PyErr) -> ErrorKind { + let s = e.instance(py) + .extract::(py) + .unwrap_or_else(|_e| "unknown exception".to_string()); + ErrorKind::CallbackError(s) +} + py_class!(class RustDirstateMap |py| { data repodir: PathBuf; data dirstate: RefCell>; @@ -172,6 +181,32 @@ .map_err(|e| PyErr::new::(py, e.description())) } + def visittrackedfiles(&self, target: PyObject) -> PyResult { + let mut dirstate = self.dirstate(py).borrow_mut(); + let mut visitor = |filepath: &Vec, _state: &FileState| { + let filename = PyBytes::new(py, &filepath.concat()).into_object(); + target.call(py, (filename,), None).map_err(|e| callback_error(py, e))?; + Ok(()) + }; + dirstate + .visit_tracked(&mut visitor) + .map_err(|e| PyErr::new::(py, e.description()))?; + Ok(py.None()) + } + + def visitremovedfiles(&self, target: PyObject) -> PyResult { + let mut dirstate = self.dirstate(py).borrow_mut(); + let mut visitor = |filepath: &Vec, _state: &FileState| { + let filename = PyBytes::new(py, &filepath.concat()).into_object(); + target.call(py, (filename,), None).map_err(|e| callback_error(py, e))?; + Ok(()) + }; + dirstate + .visit_removed(&mut visitor) + .map_err(|e| PyErr::new::(py, e.description()))?; + Ok(py.None()) + } + // Get the next dirstate object after the provided filename. If the filename is None, // returns the first file in the tree. If the provided filename is the last file, returns // None. diff --git a/rust/treedirstate/src/tree.rs b/rust/treedirstate/src/tree.rs --- a/rust/treedirstate/src/tree.rs +++ b/rust/treedirstate/src/tree.rs @@ -224,6 +224,31 @@ } } + // Visit all of the files in under this node, by calling the visitor function on each one. + fn visit<'a, F>( + &'a mut self, + store: &StoreView, + path: &mut Vec>, + visitor: &mut F, + ) -> Result<()> + where + F: FnMut(&Vec, &T) -> Result<()>, + { + for (name, entry) in self.load_entries(store)?.iter_mut() { + path.push(name); + match entry { + &mut NodeEntry::Directory(ref mut node) => { + node.visit(store, path, visitor)?; + } + &mut NodeEntry::File(ref file) => { + visitor(path, file)?; + } + } + path.pop(); + } + Ok(()) + } + /// Get the first file in the subtree under this node. If the subtree is not empty, returns a /// pair containing the path to the file as a reversed vector of key references for each path /// element, and a reference to the file. @@ -463,6 +488,14 @@ Ok(self.root.get(store, name)?) } + pub fn visit(&mut self, store: &StoreView, visitor: &mut F) -> Result<()> + where + F: FnMut(&Vec, &T) -> Result<()>, + { + let mut path = Vec::new(); + self.root.visit(store, &mut path, visitor) + } + pub fn get_first<'a>(&'a mut self, store: &StoreView) -> Result> { Ok(self.root.get_first(store)?.map(|(mut path, file)| { path.reverse(); @@ -507,7 +540,7 @@ use store::NullStore; use store::tests::MapStore; - use tree::Tree; + use tree::{KeyRef, Tree}; use filestate::FileState; // Test files in order. Note lexicographic ordering of file9 and file10. @@ -667,4 +700,25 @@ ); } + #[test] + fn visit() { + let mut ms = MapStore::new(); + let mut t = Tree::new(); + populate(&mut t, &ms); + let mut files = Vec::new(); + { + let mut v = |path: &Vec, _fs: &FileState| { + files.push(path.concat()); + Ok(()) + }; + t.visit(&mut ms, &mut v).expect("can visit"); + } + assert_eq!( + files, + TEST_FILES + .iter() + .map(|t| t.0.to_vec()) + .collect::>>() + ); + } } diff --git a/treedirstate/__init__.py b/treedirstate/__init__.py --- a/treedirstate/__init__.py +++ b/treedirstate/__init__.py @@ -18,6 +18,7 @@ ) from mercurial.i18n import _ import errno +import heapq import itertools import struct @@ -59,10 +60,9 @@ # tree each time. A future improvement is to keep the state between each # call to avoid these extra searches. class treedirstatemapiterator(object): - def __init__(self, map_, removed=False, items=False): + def __init__(self, map_, removed=False): self._rmap = map_ self._removed = removed - self._items = items self._at = None def __iter__(self): @@ -73,7 +73,7 @@ if nextitem is None: raise StopIteration self._at = nextitem[0] - return nextitem if self._items else nextitem[0] + return nextitem def next(self): return self.__next__() @@ -118,14 +118,14 @@ def itertrackeditems(self): """Returns an iterator over (filename, (state, mode, size, mtime)).""" - return treedirstatemapiterator(self._rmap, removed=False, items=True) + return treedirstatemapiterator(self._rmap, removed=False) def iterremoveditems(self): """ Returns an iterator over (filename, (state, mode, size, mtime)) for files that have been marked as removed. """ - return treedirstatemapiterator(self._rmap, removed=True, items=True) + return treedirstatemapiterator(self._rmap, removed=True) def iteritems(self): return itertools.chain(self.itertrackeditems(), @@ -162,16 +162,25 @@ return self.hastrackedfile(filename) or self.hasremovedfile(filename) def trackedfiles(self): - """Returns an iterator of all filenames tracked by the dirstate.""" - return treedirstatemapiterator(self._rmap, removed=False, items=False) + """Returns a list of all filenames tracked by the dirstate.""" + trackedfiles = [] + self._rmap.visittrackedfiles(trackedfiles.append) + return iter(trackedfiles) def removedfiles(self): - """Returns an iterator of all removed files in the dirstate.""" - return treedirstatemapiterator(self._rmap, removed=True, items=False) + """Returns a list of all removed files in the dirstate.""" + removedfiles = [] + self._rmap.visitremovedfiles(removedfiles.append) + return removedfiles def __iter__(self): """Returns an iterator of all files in the dirstate.""" - return itertools.chain(self.trackedfiles(), self.removedfiles()) + trackedfiles = self.trackedfiles() + removedfiles = self.removedfiles() + if removedfiles: + return heapq.merge(iter(trackedfiles), iter(removedfiles)) + else: + return iter(trackedfiles) def keys(self): return list(iter(self))