diff --git a/rustworkx/rustworkx.pyi b/rustworkx/rustworkx.pyi index ed56547d7..34c3ff536 100644 --- a/rustworkx/rustworkx.pyi +++ b/rustworkx/rustworkx.pyi @@ -980,6 +980,7 @@ class _RustworkxCustomVecIter(Generic[_T_co], Sequence[_T_co], ABC): def __setstate__(self, state: Sequence[_T_co]) -> None: ... def __array__(self, _dt: np.dtype | None = ...) -> np.ndarray: ... def __iter__(self) -> Iterator[_T_co]: ... + def __reversed__(self) -> Iterator[_T_co]: ... class _RustworkxCustomHashMapIter(Generic[_S, _T_co], Mapping[_S, _T_co], ABC): def __init__(self) -> None: ... diff --git a/src/iterators.rs b/src/iterators.rs index 9fdf9d2ec..ace324cb6 100644 --- a/src/iterators.rs +++ b/src/iterators.rs @@ -14,8 +14,8 @@ // // :`custom_vec_iter_impl` holds a `Vec` and can be used as a // read-only sequence/list. To use it, you should specify the name of the new type for the -// iterable, a name for that new type's iterator, the name of the vector that holds the data, the -// type `T` and a docstring. +// iterable, a name for that new type's iterator, a name for the new type's reversed iterator, the +// name of the vector that holds the data, the type `T` and a docstring. // // e.g `custom_vec_iter_impl!(MyReadOnlyType, data, (usize, f64), "Docs");` // defines a new type named `MyReadOnlyType` that holds a vector called `data` @@ -474,7 +474,7 @@ impl PyConvertToPyArray for Vec<(usize, usize, PyObject)> { } macro_rules! custom_vec_iter_impl { - ($name:ident, $iter:ident, $data:ident, $T:ty, $doc:literal) => { + ($name:ident, $iter:ident, $reversed:ident, $data:ident, $T:ty, $doc:literal) => { #[doc = $doc] #[pyclass(module = "rustworkx", sequence)] #[derive(Clone)] @@ -588,6 +588,13 @@ macro_rules! custom_vec_iter_impl { } } + fn __reversed__(self_: Py, py: Python) -> $reversed { + $reversed { + inner: Some(self_.clone_ref(py)), + index: 0, + } + } + fn __array__(&self, py: Python, _dt: Option<&PyArrayDescr>) -> PyResult { // Note: we accept the dtype argument on the signature but // effictively do nothing with it to let Numpy handle the conversion itself @@ -652,12 +659,64 @@ macro_rules! custom_vec_iter_impl { self.inner = None; } } + + #[doc = concat!("Custom reversed iterator class for :class:`.", stringify!($name), "`")] + // No module because this isn't constructable from Python space, and is only exposed as an + // implementation detail. + #[pyclass] + pub struct $reversed { + inner: Option>, + index: usize, + } + + #[pymethods] + impl $reversed { + fn __next__(&mut self, py: Python) -> Option> { + let data = self.inner.as_ref().unwrap().borrow(py); + let len = data.$data.len(); + if self.index < len { + let out = data.$data[len - self.index - 1].clone().into_py(py); + self.index += 1; + Some(out) + } else { + None + } + } + + fn __iter__(self_: Py) -> Py { + // Python iterators typically just return themselves from this, though in principle + // we could return a separate object that iterates starting from the same point. + self_ + } + + fn __length_hint__(&self, py: Python) -> usize { + self.inner + .as_ref() + .unwrap() + .borrow(py) + .$data + .len() + .saturating_sub(self.index) + } + + fn __traverse__(&self, vis: PyVisit) -> Result<(), PyTraverseError> { + if let Some(obj) = self.inner.as_ref() { + vis.call(obj)? + } + Ok(()) + } + + fn __clear__(&mut self) { + self.inner = None; + } + } }; } custom_vec_iter_impl!( BFSSuccessors, BFSSuccessorsIter, + BFSSuccessorsRev, bfs_successors, (PyObject, Vec), "A custom class for the return from :func:`rustworkx.bfs_successors` @@ -711,6 +770,7 @@ impl PyGCProtocol for BFSSuccessors { custom_vec_iter_impl!( BFSPredecessors, BFSPredecessorsIter, + BFSPredecessorsRev, bfs_predecessors, (PyObject, Vec), "A custom class for the return from :func:`rustworkx.bfs_predecessors` @@ -764,6 +824,7 @@ impl PyGCProtocol for BFSPredecessors { custom_vec_iter_impl!( NodeIndices, NodeIndicesIter, + NodeIndicesRev, nodes, usize, "A custom class for the return of node indices @@ -797,6 +858,7 @@ impl PyGCProtocol for NodeIndices {} custom_vec_iter_impl!( EdgeList, EdgeListIter, + EdgeListRev, edges, (usize, usize), "A custom class for the return of edge lists @@ -836,6 +898,7 @@ impl PyGCProtocol for EdgeList {} custom_vec_iter_impl!( WeightedEdgeList, WeightedEdgeListIter, + WeightedEdgeListRev, edges, (usize, usize, PyObject), "A custom class for the return of edge lists with weights @@ -887,6 +950,7 @@ impl PyGCProtocol for WeightedEdgeList { custom_vec_iter_impl!( EdgeIndices, EdgeIndicesIter, + EdgeIndicesRev, edges, usize, "A custom class for the return of edge indices @@ -940,6 +1004,7 @@ impl PyDisplay for EdgeList { custom_vec_iter_impl!( Chains, ChainsIter, + ChainsRev, chains, EdgeList, "A custom class for the return of a list of list of edges.