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 @@ -18,7 +18,7 @@ exceptions::GraphError, }; use cpython::{ - ObjectProtocol, PyDict, PyModule, PyObject, PyResult, Python, + ObjectProtocol, PyDict, PyModule, PyObject, PyResult, PyTuple, Python, PythonObject, ToPyObject, }; use hg::discovery::PartialDiscovery as CorePartialDiscovery; @@ -111,6 +111,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