Solidify Batch API

4 files changed, 98 insertions, 25 deletions

diff --git a/example.py b/example.py
new file mode 100644
index 0000000..08fd8b0
--- /dev/null
+++ b/example.py
@@ -0,0 +1,43 @@
+import asyncio
+
+from rivulet import Pipeline, Batch
+
+
+async def source():
+    for i in range(22):
+        yield i
+
+
+async def double(gen):
+    async for value in gen:
+        yield value * 2
+
+
+async def dump(gen):
+    async for x in gen:
+        print(x)
+        yield x
+
+
+async def sum(batches):
+    async for batch in batches:
+        res = 0
+        for x in batch:
+            res += x
+        yield res
+
+
+async def main():
+    pipe = Pipeline(source())
+    pipe.add_step(double)
+
+    batch = Batch(N=5, timeout=0.1)
+    pipe.add_step(batch)
+    pipe.add_step(dump)
+    pipe.add_step(sum)
+
+    res = await pipe.collect()
+    print(res)
+
+
+asyncio.run(main())
diff --git a/rivulet/__init__.py b/rivulet/__init__.py
index bbc3523..3b42d19 100644
--- a/rivulet/__init__.py
+++ b/rivulet/__init__.py
@@ -1,3 +1,4 @@
-from .batch import BatchProcessor
+from .batch import Batch
+from .pipeline import Pipeline
 
-__all__ = [BatchProcessor]
+__all__ = [Batch, Pipeline]
diff --git a/rivulet/batch.py b/rivulet/batch.py
index 57266b8..9a6e863 100644
--- a/rivulet/batch.py
+++ b/rivulet/batch.py
@@ -1,17 +1,50 @@
 import time
 
-from typing import TypeVar, Generic, AsyncIterable, AsyncIterator, List
-
+from typing import AsyncGenerator, AsyncIterable, List, TypeVar, Generic
 
 T = TypeVar("T")
+U = TypeVar("U")
+
+
+class Batch(Generic[T]):
+    """
+    A processor that batches items from an async generator based on batch size or timeout.
+
+    This can be used as a step in a Pipeline to batch items from the previous step.
+    """
+
+    def __init__(self, N: int, timeout: float):
+        self.N = N
+        self.timeout = timeout
+
+    def __call__(
+        self, source: AsyncGenerator[T, None]
+    ) -> AsyncGenerator[List[T], None]:
+        """
+        Make the BatchProcessor callable so it can be added as a step in the Pipeline.
+
+        Args:
+            source: The source async generator providing individual items
+
+        Returns:
+            An async generator yielding batches of items as lists
+        """
+        return self.process(source)
+
+    async def process(self, source: AsyncIterable[T]) -> AsyncGenerator[List[T], None]:
+        """
+        Process items from the source generator into batches.
 
+        Batches are yielded when either:
+        1. The batch size is reached
+        2. The timeout period has elapsed and there are items in the buffer
 
-class BatchProcessor(Generic[T]):
-    def __init__(self, batch_size: int, timeout_seconds: float):
-        self.batch_size = batch_size
-        self.timeout_seconds = timeout_seconds
+        Args:
+            source: The source async generator providing individual items
 
-    async def process(self, source: AsyncIterable[T]) -> AsyncIterator[List[T]]:
+        Yields:
+            Lists containing batches of items from the source
+        """
         buffer: List[T] = []
         last_flush_time = time.time()
 
@@ -19,12 +52,12 @@ class BatchProcessor(Generic[T]):
             buffer.append(item)
 
             current_time = time.time()
-            timeout_reached = current_time - last_flush_time >= self.timeout_seconds
-            buffer_full = len(buffer) >= self.batch_size
+            timeout_reached = current_time - last_flush_time >= self.timeout
+            buffer_full = len(buffer) >= self.N
 
             if buffer_full or (timeout_reached and buffer):
-                yield buffer
-                buffer = []
+                yield buffer.copy()  # Yield a copy to avoid mutation issues
+                buffer.clear()
                 last_flush_time = current_time
 
         # Don't forget items in buffer when source is exhausted
diff --git a/rivulet/tests/test_batch.py b/rivulet/tests/test_batch.py
index fe8be9a..6b32462 100644
--- a/rivulet/tests/test_batch.py
+++ b/rivulet/tests/test_batch.py
@@ -1,16 +1,14 @@
 import asyncio
 import pytest
 
-from ..batch import BatchProcessor  # Update with your actual import
+from ..batch import Batch  # Update with your actual import
 
 
-class TestBatchProcessor:
+class TestBatch:
     @pytest.mark.asyncio
     async def test_batch_by_size(self):
         # Test batching by size
-        batch_processor = BatchProcessor[int](
-            batch_size=3, timeout_seconds=10.0
-        )  # Long timeout
+        batch_processor = Batch[int](N=3, timeout=10.0)  # Long timeout
 
         async def source():
             for i in range(8):  # 8 items should produce 2 full batches and 1 partial
@@ -28,9 +26,7 @@ class TestBatchProcessor:
     @pytest.mark.asyncio
     async def test_batch_by_timeout(self):
         # Test batching by timeout
-        batch_processor = BatchProcessor[int](
-            batch_size=10, timeout_seconds=0.2
-        )  # Small timeout
+        batch_processor = Batch[int](N=10, timeout=0.2)  # Small timeout
 
         async def slow_source():
             for i in range(5):
@@ -49,7 +45,7 @@ class TestBatchProcessor:
     @pytest.mark.asyncio
     async def test_empty_source(self):
         # Test with empty source
-        batch_processor = BatchProcessor[int](batch_size=3, timeout_seconds=0.5)
+        batch_processor = Batch[int](N=3, timeout=0.5)
 
         async def empty_source():
             if False:  # Never yields
@@ -64,7 +60,7 @@ class TestBatchProcessor:
     @pytest.mark.asyncio
     async def test_exact_batch_size(self):
         # Test with source that produces exactly one full batch
-        batch_processor = BatchProcessor[int](batch_size=3, timeout_seconds=0.5)
+        batch_processor = Batch[int](N=3, timeout=0.5)
 
         async def exact_source():
             for i in range(3):
@@ -84,7 +80,7 @@ class TestBatchProcessor:
             def __init__(self, value):
                 self.value = value
 
-        batch_processor = BatchProcessor[TestItem](batch_size=2, timeout_seconds=0.5)
+        batch_processor = Batch[TestItem](N=2, timeout=0.5)
 
         async def object_source():
             for i in range(3):
@@ -101,7 +97,7 @@ class TestBatchProcessor:
     @pytest.mark.asyncio
     async def test_concurrent_items(self):
         # Test with items arriving close together but processed in batches
-        batch_processor = BatchProcessor[int](batch_size=5, timeout_seconds=0.3)
+        batch_processor = Batch[int](N=5, timeout=0.3)
 
         async def concurrent_source():
             # Produce items quickly