use crate::conversion::{py_set, rev_pyiter_collect};
use cindex::Index;
use cpython::{
ObjectProtocol, PyClone, PyDict, PyList, PyModule, PyObject, PyResult,
Python, PythonObject, ToPyObject,
};
use exceptions::GraphError;
use hg::Revision;
use hg::{
AncestorsIterator as CoreIterator, LazyAncestors as CoreLazy,
MissingAncestors as CoreMissing,
};
use std::cell::RefCell;
use std::collections::HashSet;
py_class!(pub class AncestorsIterator |py| {
data inner: RefCell<Box<CoreIterator<Index>>>;
def __next__(&self) -> PyResult<Option<Revision>> {
match self.inner(py).borrow_mut().next() {
Some(Err(e)) => Err(GraphError::pynew(py, e)),
None => Ok(None),
Some(Ok(r)) => Ok(Some(r)),
}
}
def __contains__(&self, rev: Revision) -> PyResult<bool> {
self.inner(py).borrow_mut().contains(rev)
.map_err(|e| GraphError::pynew(py, e))
}
def __iter__(&self) -> PyResult<Self> {
Ok(self.clone_ref(py))
}
def __new__(_cls, index: PyObject, initrevs: PyObject, stoprev: Revision,
inclusive: bool) -> PyResult<AncestorsIterator> {
let initvec: Vec<Revision> = rev_pyiter_collect(py, &initrevs)?;
let ait = CoreIterator::new(
Index::new(py, index)?,
initvec,
stoprev,
inclusive,
)
.map_err(|e| GraphError::pynew(py, e))?;
AncestorsIterator::from_inner(py, ait)
}
});
impl AncestorsIterator {
pub fn from_inner(py: Python, ait: CoreIterator<Index>) -> PyResult<Self> {
Self::create_instance(py, RefCell::new(Box::new(ait)))
}
}
py_class!(pub class LazyAncestors |py| {
data inner: RefCell<Box<CoreLazy<Index>>>;
def __contains__(&self, rev: Revision) -> PyResult<bool> {
self.inner(py)
.borrow_mut()
.contains(rev)
.map_err(|e| GraphError::pynew(py, e))
}
def __iter__(&self) -> PyResult<AncestorsIterator> {
AncestorsIterator::from_inner(py, self.inner(py).borrow().iter())
}
def __bool__(&self) -> PyResult<bool> {
Ok(!self.inner(py).borrow().is_empty())
}
def __new__(_cls, index: PyObject, initrevs: PyObject, stoprev: Revision,
inclusive: bool) -> PyResult<Self> {
let initvec: Vec<Revision> = rev_pyiter_collect(py, &initrevs)?;
let lazy =
CoreLazy::new(Index::new(py, index)?, initvec, stoprev, inclusive)
.map_err(|e| GraphError::pynew(py, e))?;
Self::create_instance(py, RefCell::new(Box::new(lazy)))
}
});
py_class!(pub class MissingAncestors |py| {
data inner: RefCell<Box<CoreMissing<Index>>>;
def __new__(_cls, index: PyObject, bases: PyObject) -> PyResult<MissingAncestors> {
let bases_vec: Vec<Revision> = rev_pyiter_collect(py, &bases)?;
let inner = CoreMissing::new(Index::new(py, index)?, bases_vec);
MissingAncestors::create_instance(py, RefCell::new(Box::new(inner)))
}
def hasbases(&self) -> PyResult<bool> {
Ok(self.inner(py).borrow().has_bases())
}
def addbases(&self, bases: PyObject) -> PyResult<PyObject> {
let mut inner = self.inner(py).borrow_mut();
let bases_vec: Vec<Revision> = rev_pyiter_collect(py, &bases)?;
inner.add_bases(bases_vec);
Ok(py.None())
}
def bases(&self) -> PyResult<PyObject> {
py_set(py, self.inner(py).borrow().get_bases())
}
def basesheads(&self) -> PyResult<PyObject> {
let inner = self.inner(py).borrow();
py_set(py, &inner.bases_heads().map_err(|e| GraphError::pynew(py, e))?)
}
def removeancestorsfrom(&self, revs: PyObject) -> PyResult<PyObject> {
let mut inner = self.inner(py).borrow_mut();
let mut revs_pyset: HashSet<Revision> = rev_pyiter_collect(py, &revs)?;
inner.remove_ancestors_from(&mut revs_pyset)
.map_err(|e| GraphError::pynew(py, e))?;
let mut remaining_pyint_vec: Vec<PyObject> = Vec::with_capacity(
revs_pyset.len());
for rev in revs_pyset {
remaining_pyint_vec.push(rev.to_py_object(py).into_object());
}
let remaining_pylist = PyList::new(py, remaining_pyint_vec.as_slice());
revs.call_method(py, "intersection_update", (remaining_pylist, ), None)
}
def missingancestors(&self, revs: PyObject) -> PyResult<PyList> {
let mut inner = self.inner(py).borrow_mut();
let revs_vec: Vec<Revision> = rev_pyiter_collect(py, &revs)?;
let missing_vec = match inner.missing_ancestors(revs_vec) {
Ok(missing) => missing,
Err(e) => {
return Err(GraphError::pynew(py, e));
}
};
let mut missing_pyint_vec: Vec<PyObject> = Vec::with_capacity(
missing_vec.len());
for rev in missing_vec {
missing_pyint_vec.push(rev.to_py_object(py).into_object());
}
Ok(PyList::new(py, missing_pyint_vec.as_slice()))
}
});
pub fn init_module(py: Python, package: &str) -> PyResult<PyModule> {
let dotted_name = &format!("{}.ancestor", package);
let m = PyModule::new(py, dotted_name)?;
m.add(py, "__package__", package)?;
m.add(
py,
"__doc__",
"Generic DAG ancestor algorithms - Rust implementation",
)?;
m.add_class::<AncestorsIterator>(py)?;
m.add_class::<LazyAncestors>(py)?;
m.add_class::<MissingAncestors>(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)?;
Ok(m)
}