From a20a9e9650a87f8c3f2fc6bb8ae1b8ffdb5d0f98 Mon Sep 17 00:00:00 2001
From: "Andreas Koepf (aider)" <andreas.koepf@aleph-alpha.com>
Date: Fri, 24 Jan 2025 10:16:27 +0100
Subject: [PATCH] feat: Add return type annotation to
 ProceduralDataset.__next__()

---
 reasoning_gym/dataset.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/reasoning_gym/dataset.py b/reasoning_gym/dataset.py
index b775874c..cce6752d 100644
--- a/reasoning_gym/dataset.py
+++ b/reasoning_gym/dataset.py
@@ -22,7 +22,7 @@ class ProceduralDataset(ABC, Sized, Iterable[Dict[str, Any]]):
         self._current_idx = 0
         return self
     
-    def __next__(self):
+    def __next__(self) -> Dict[str, Any]:
         """Get next item in iteration"""
         if self._current_idx >= self.size:
             raise StopIteration