diff --git a/rust/hg-cpython/src/discovery.rs b/rust/hg-cpython/src/discovery.rs --- a/rust/hg-cpython/src/discovery.rs +++ b/rust/hg-cpython/src/discovery.rs @@ -15,7 +15,8 @@ use crate::conversion::{py_set, rev_pyiter_collect}; use cindex::Index; use cpython::{ - ObjectProtocol, PyDict, PyModule, PyObject, PyResult, Python, ToPyObject, + ObjectProtocol, PyDict, PyModule, PyObject, PyResult, PyTuple, Python, + PythonObject, ToPyObject, }; use exceptions::GraphError; use hg::discovery::PartialDiscovery as CorePartialDiscovery; @@ -101,6 +102,32 @@ .map_err(|e| GraphError::pynew(py, e))? ) } + + def takefullsample(&self, _headrevs: PyObject, + size: usize) -> PyResult { + let mut inner = self.inner(py).borrow_mut(); + let sample = inner.take_full_sample(size) + .map_err(|e| GraphError::pynew(py, e))?; + let as_vec: Vec = sample + .iter() + .map(|rev| rev.to_py_object(py).into_object()) + .collect(); + Ok(PyTuple::new(py, as_vec.as_slice()).into_object()) + } + + def takequicksample(&self, headrevs: PyObject, + size: usize) -> PyResult { + let mut inner = self.inner(py).borrow_mut(); + let revsvec: Vec = rev_pyiter_collect(py, &headrevs)?; + let sample = inner.take_quick_sample(revsvec, size) + .map_err(|e| GraphError::pynew(py, e))?; + let as_vec: Vec = sample + .iter() + .map(|rev| rev.to_py_object(py).into_object()) + .collect(); + Ok(PyTuple::new(py, as_vec.as_slice()).into_object()) + } + }); /// Create the module, with __package__ given from parent