From 3a8a613573c8e1ad53e27f811e637fba557d165d Mon Sep 17 00:00:00 2001
From: "Andreas Koepf (aider)" <andreas.koepf@aleph-alpha.com>
Date: Fri, 24 Jan 2025 10:14:42 +0100
Subject: [PATCH] feat: Add Sized and Iterable base classes to
 ProceduralDataset

---
 reasoning_gym/dataset.py | 5 +++--
 1 file changed, 3 insertions(+), 2 deletions(-)

diff --git a/reasoning_gym/dataset.py b/reasoning_gym/dataset.py
index 8cca17ff..b775874c 100644
--- a/reasoning_gym/dataset.py
+++ b/reasoning_gym/dataset.py
@@ -1,10 +1,11 @@
 """Base class for procedural dataset generators"""
 from abc import ABC, abstractmethod
+from collections.abc import Sized, Iterable
 from random import Random
-from typing import Optional
+from typing import Optional, Iterator, Dict, Any
 
 
-class ProceduralDataset(ABC):
+class ProceduralDataset(ABC, Sized, Iterable[Dict[str, Any]]):
     """Abstract base class for procedural dataset generators"""
     
     def __init__(self, seed: Optional[int] = None, size: int = 500):