From a4ec6d0aa7b9f43866fbb91326595972e732253a Mon Sep 17 00:00:00 2001 From: Yuval Adam <_@yuv.al> Date: Fri, 21 Mar 2025 22:51:42 +0100 Subject: Add step init --- README.md | 4 ++-- rivulet/pipeline.py | 4 +++- rivulet/tests/test_pipeline.py | 18 ++++++++++++++++++ 3 files changed, 23 insertions(+), 3 deletions(-) diff --git a/README.md b/README.md index 08becf1..9b9d744 100644 --- a/README.md +++ b/README.md @@ -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 @@ -78,6 +78,24 @@ async def test_empty_pipeline(): assert results == ["test"] +@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""" -- cgit v1.3.1