diff options
| -rw-r--r-- | README.md | 4 | ||||
| -rw-r--r-- | rivulet/pipeline.py | 4 | ||||
| -rw-r--r-- | rivulet/tests/test_pipeline.py | 18 |
3 files changed, 23 insertions, 3 deletions
@@ -27,8 +27,8 @@ async def main(): async for out in pipe: print(out) - # or collect them all - res = await pipe.collect() + # or just single line it and collect them all + res = await Pipeline(source(), double, batch, sum).collect() ``` ## License diff --git a/rivulet/pipeline.py b/rivulet/pipeline.py index 6efd821..097b1a0 100644 --- a/rivulet/pipeline.py +++ b/rivulet/pipeline.py @@ -12,12 +12,14 @@ class Pipeline: Each step is a function that takes an async generator and returns a new async generator. """ - def __init__(self, source: AsyncGenerator[Any, None]): + def __init__(self, source: AsyncGenerator[Any, None], *steps): """Initialize the pipeline with a source async generator.""" self.source = source self.steps: List[ Callable[[AsyncGenerator[Any, None]], AsyncGenerator[Any, None]] ] = [] + for step in steps: + self.add_step(step) def add_step( self, diff --git a/rivulet/tests/test_pipeline.py b/rivulet/tests/test_pipeline.py index 33462bd..42ac055 100644 --- a/rivulet/tests/test_pipeline.py +++ b/rivulet/tests/test_pipeline.py @@ -79,6 +79,24 @@ async def test_empty_pipeline(): @pytest.mark.asyncio +async def test_pipeline_init_steps(): + """Test pipeline with steps init in constructor""" + + async def source(): + yield "test" + + async def dupe(gen): + async for value in gen: + for _ in range(3): + yield value + + pipeline = Pipeline(source(), dupe, dupe) + results = await pipeline.collect() + + assert results == ["test"] * 9 + + +@pytest.mark.asyncio async def test_pipeline_with_filtering(): """Test pipeline with a transformation that filters values""" |