feat: implement iteration for the WorkingSet

mypy.ini

@@ -5,4 +5,4 @@ no_implicit_optional = True
 check_untyped_defs = True
 warn_unused_ignores = True
 show_error_codes = True
-disable_error_code = assignment 
+disable_error_code = assignment

src/working_set.rs

@@ -5,6 +5,21 @@ use taskchampion::WorkingSet as TCWorkingSet;
 #[pyclass]
 pub struct WorkingSet(pub(crate) TCWorkingSet);
 
+#[pyclass]
+struct WorkingSetIter {
+    iter: std::vec::IntoIter<(usize, String)>,
+}
+
+#[pymethods]
+impl WorkingSetIter {
+    fn __iter__(slf: PyRef<'_, Self>) -> PyRef<'_, Self> {
+        slf
+    }
+
+    fn __next__(mut slf: PyRefMut<'_, Self>) -> Option<(usize, String)> {
+        slf.iter.next()
+    }
+}
 #[pymethods]
 impl WorkingSet {
     pub fn __len__(&self) -> usize {
@@ -28,9 +43,15 @@ impl WorkingSet {
         self.0.by_uuid(Uuid::parse_str(&uuid).unwrap())
     }
 
-    fn __iter__(_slf: PyRef<'_, Self>) -> PyResult<Py<WorkingSet>> {
-        todo!("Figure way to propertly implement iterator for python")
-        // Usability-wise we want it to hold the reference to the iterator, so that
-        // with each iteration the state persists.
+    fn __iter__(slf: PyRef<'_, Self>) -> PyResult<Py<WorkingSetIter>> {
+        let iter = slf
+            .0
+            .iter()
+            .map(|(i, id)| (i, id.to_string()))
+            .collect::<Vec<_>>()
+            .into_iter();
+        let iter = WorkingSetIter { iter };
+
+        Py::new(slf.py(), iter)
     }
 }

taskchampion.pyi

@@ -1,6 +1,6 @@
-from typing import Optional
 from datetime import datetime
 from enum import Enum
+from typing import Optional, Iterator
 
 
 class Replica:
@@ -110,6 +110,8 @@ class WorkingSet:
     def is_empty(self) -> bool: ...
     def by_index(self, index: int) -> Optional[str]: ...
     def by_uuid(self, uuid: str) -> Optional[int]: ...
+    def __iter__(self) -> Iterator[tuple[int, str]]: ...
+    def __next__(self) -> tuple[int, str]: ...
 
 
 class Annotation:

tests/test_working_set.py

@@ -1,4 +1,4 @@
-from taskchampion import Replica, WorkingSet, Status, Operation
+from taskchampion import Replica, WorkingSet, Status
 from pathlib import Path
 import pytest
 import uuid
@@ -39,14 +39,14 @@ def test_by_index(working_set: WorkingSet):
     assert working_set.by_index(1) is not None
 
 
-@pytest.mark.skip()
 def test_iter(working_set: WorkingSet):
     assert iter(working_set)
 
 
-@pytest.mark.skip()
 def test_next(working_set: WorkingSet):
-    assert next(working_set)[0] == 1
-    assert next(working_set)[0] == 2
-    with pytest.raises(OSError):
-        next(working_set)
+    working_set_iterator = iter(working_set)
+
+    assert next(working_set_iterator)[0] == 1
+    assert next(working_set_iterator)[0] == 2
+    with pytest.raises(StopIteration):
+        next(working_set_iterator)