From fc84959668c855b6f10647f28f455d68ab777d12 Mon Sep 17 00:00:00 2001
From: Illya Laifu <illyalaifu@gmail.com>
Date: Thu, 12 Sep 2024 00:53:28 +0300
Subject: [PATCH] feat: implement iteration for the WorkingSet

---
 mypy.ini                  |  2 +-
 src/working_set.rs        | 29 +++++++++++++++++++++++++----
 taskchampion.pyi          |  4 +++-
 tests/test_working_set.py | 14 +++++++-------
 4 files changed, 36 insertions(+), 13 deletions(-)

diff --git a/mypy.ini b/mypy.ini
index 1d270f3..f742e93 100644
--- a/mypy.ini
+++ b/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
diff --git a/src/working_set.rs b/src/working_set.rs
index 720be0e..eccef6a 100644
--- a/src/working_set.rs
+++ b/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)
     }
 }
diff --git a/taskchampion.pyi b/taskchampion.pyi
index 8041327..4211dd4 100644
--- a/taskchampion.pyi
+++ b/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:
diff --git a/tests/test_working_set.py b/tests/test_working_set.py
index e8fbd9f..d5c76e5 100644
--- a/tests/test_working_set.py
+++ b/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)