Add Pipeline

2 files changed, 161 insertions, 0 deletions

diff --git a/rivulet/pipeline.py b/rivulet/pipeline.py
new file mode 100644
index 0000000..226a645
--- /dev/null
+++ b/rivulet/pipeline.py
@@ -0,0 +1,63 @@
+from typing import AsyncGenerator, TypeVar, Callable, List, Generic, Any
+
+T = TypeVar('T')
+U = TypeVar('U')
+V = TypeVar('V')
+
+class Pipeline:
+    """
+    A flexible pipeline that chains multiple async generators.
+
+    Each step is a function that takes an async generator and returns a new async generator.
+    """
+
+    def __init__(self, source: AsyncGenerator[Any, None]):
+        """Initialize the pipeline with a source async generator."""
+        self.source = source
+        self.steps: List[Callable[[AsyncGenerator[Any, None]], AsyncGenerator[Any, None]]] = []
+
+    def add_step(self, transform: Callable[[AsyncGenerator[Any, None]], AsyncGenerator[Any, None]]):
+        """
+        Add a transformation step to the pipeline.
+
+        Args:
+            transform: A function that takes an async generator and returns an async generator
+
+        Returns:
+            The pipeline instance for method chaining
+        """
+        self.steps.append(transform)
+        return self
+
+    def __aiter__(self):
+        """Make the pipeline itself an async generator."""
+        return self._execute()
+
+    async def _execute(self):
+        """
+        Execute the pipeline by chaining all generators together.
+
+        Yields:
+            Values from the final step of the pipeline
+        """
+        current_gen = self.source
+
+        # Apply each transformation step
+        for step in self.steps:
+            current_gen = step(current_gen)
+
+        # Yield all items from the final generator
+        async for item in current_gen:
+            yield item
+
+    async def collect(self):
+        """
+        Collect all results from the pipeline into a list.
+
+        Returns:
+            A list containing all output items from the pipeline
+        """
+        results = []
+        async for item in self:
+            results.append(item)
+        return results
\ No newline at end of file
diff --git a/rivulet/tests/test_pipeline.py b/rivulet/tests/test_pipeline.py
new file mode 100644
index 0000000..33462bd
--- /dev/null
+++ b/rivulet/tests/test_pipeline.py
@@ -0,0 +1,98 @@
+import pytest
+
+from ..pipeline import Pipeline
+
+
+@pytest.mark.asyncio
+async def test_pipeline_basic_transformations():
+    """Test pipeline with simple transform functions that process values"""
+
+    async def source():
+        for i in range(3):
+            yield i
+
+    async def double(gen):
+        async for value in gen:
+            yield value * 2
+
+    async def as_string(gen):
+        async for value in gen:
+            yield str(value)
+
+    # Create and execute pipeline
+    pipeline = Pipeline(source())
+    pipeline.add_step(double)
+    pipeline.add_step(as_string)
+
+    # Collect results
+    results = await pipeline.collect()
+
+    # Verify results
+    assert results == ["0", "2", "4"]
+
+
+@pytest.mark.asyncio
+async def test_pipeline_with_expanding_transformations():
+    """Test pipeline with a transformation that outputs multiple values per input"""
+
+    async def source():
+        yield 1
+        yield 2
+
+    async def duplicate(gen):
+        async for value in gen:
+            yield value
+            yield value
+
+    async def multiply_by_ten(gen):
+        async for value in gen:
+            yield value * 10
+
+    # Test pipeline with expanding transformation first
+    pipeline1 = Pipeline(source())
+    pipeline1.add_step(duplicate)
+    pipeline1.add_step(multiply_by_ten)
+
+    results1 = await pipeline1.collect()
+    assert results1 == [10, 10, 20, 20]
+
+    # Test pipeline with expanding transformation last
+    pipeline2 = Pipeline(source())
+    pipeline2.add_step(multiply_by_ten)
+    pipeline2.add_step(duplicate)
+
+    results2 = await pipeline2.collect()
+    assert results2 == [10, 10, 20, 20]
+
+
+@pytest.mark.asyncio
+async def test_empty_pipeline():
+    """Test pipeline with no transformations"""
+
+    async def source():
+        yield "test"
+
+    pipeline = Pipeline(source())
+    results = await pipeline.collect()
+
+    assert results == ["test"]
+
+
+@pytest.mark.asyncio
+async def test_pipeline_with_filtering():
+    """Test pipeline with a transformation that filters values"""
+
+    async def source():
+        for i in range(5):
+            yield i
+
+    async def even_only(gen):
+        async for value in gen:
+            if value % 2 == 0:
+                yield value
+
+    pipeline = Pipeline(source())
+    pipeline.add_step(even_only)
+
+    results = await pipeline.collect()
+    assert results == [0, 2, 4]