example.py


1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
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())