diff --git a/rust/hg-core/src/dirstate/dirstate_map.rs b/rust/hg-core/src/dirstate/dirstate_map.rs --- a/rust/hg-core/src/dirstate/dirstate_map.rs +++ b/rust/hg-core/src/dirstate/dirstate_map.rs @@ -237,6 +237,25 @@ ) } + /// Useful to get immutable references to those sets in contexts where + /// you only have an immutable reference to the `DirstateMap`, like when + /// sharing references with Python. + /// + /// TODO, get rid of this along with the other "setter/getter" stuff when + /// a nice typestate plan is defined. + /// + /// # Panics + /// + /// Will panic if either set is `None`. + pub fn get_non_normal_other_parent_entries_panic( + &self, + ) -> (&HashSet, &HashSet) { + ( + self.non_normal_set.as_ref().unwrap(), + self.other_parent_set.as_ref().unwrap(), + ) + } + pub fn set_non_normal_other_parent_entries(&mut self, force: bool) { if !force && self.non_normal_set.is_some() diff --git a/rust/hg-cpython/src/dirstate/dirstate_map.rs b/rust/hg-cpython/src/dirstate/dirstate_map.rs --- a/rust/hg-cpython/src/dirstate/dirstate_map.rs +++ b/rust/hg-cpython/src/dirstate/dirstate_map.rs @@ -20,7 +20,9 @@ use crate::{ dirstate::copymap::{CopyMap, CopyMapItemsIterator, CopyMapKeysIterator}, - dirstate::non_normal_entries::NonNormalEntries, + dirstate::non_normal_entries::{ + NonNormalEntries, NonNormalEntriesIterator, + }, dirstate::{dirs_multiset::Dirs, make_dirstate_tuple}, }; use hg::{ @@ -244,6 +246,22 @@ Ok(ret) } + def non_normal_entries_iter(&self) -> PyResult { + // Make sure the sets are defined before we no longer have a mutable + // reference to the dmap. + self.inner(py) + .borrow_mut() + .set_non_normal_other_parent_entries(false); + + let leaked_ref = self.inner(py).leak_immutable(); + + NonNormalEntriesIterator::from_inner(py, unsafe { + leaked_ref.map(py, |o| { + o.get_non_normal_other_parent_entries_panic().0.iter() + }) + }) + } + def hastrackeddir(&self, d: PyObject) -> PyResult { let d = d.extract::(py)?; Ok(self.inner(py).borrow_mut() diff --git a/rust/hg-cpython/src/dirstate/non_normal_entries.rs b/rust/hg-cpython/src/dirstate/non_normal_entries.rs --- a/rust/hg-cpython/src/dirstate/non_normal_entries.rs +++ b/rust/hg-cpython/src/dirstate/non_normal_entries.rs @@ -6,11 +6,15 @@ // GNU General Public License version 2 or any later version. use cpython::{ - exc::NotImplementedError, CompareOp, ObjectProtocol, PyErr, PyList, - PyObject, PyResult, PyString, Python, PythonObject, ToPyObject, + exc::NotImplementedError, CompareOp, ObjectProtocol, PyBytes, PyClone, + PyErr, PyList, PyObject, PyResult, PyString, Python, PythonObject, + ToPyObject, UnsafePyLeaked, }; use crate::dirstate::DirstateMap; +use hg::utils::hg_path::HgPathBuf; +use std::cell::RefCell; +use std::collections::hash_set; py_class!(pub class NonNormalEntries |py| { data dmap: DirstateMap; @@ -34,6 +38,10 @@ def __repr__(&self) -> PyResult { self.dmap(py).non_normal_entries_display(py) } + + def __iter__(&self) -> PyResult { + self.dmap(py).non_normal_entries_iter(py) + } }); impl NonNormalEntries { @@ -49,4 +57,20 @@ } Ok(true) } + + fn translate_key( + py: Python, + key: &HgPathBuf, + ) -> PyResult> { + Ok(Some(PyBytes::new(py, key.as_ref()))) + } } + +type NonNormalEntriesIter<'a> = hash_set::Iter<'a, HgPathBuf>; + +py_shared_iterator!( + NonNormalEntriesIterator, + UnsafePyLeaked>, + NonNormalEntries::translate_key, + Option +);