merge: Ability to identify source of an "event"

[madduck]

When I merge two streams like so:

sources = aiostream.stream.merge(genstrings1, genlists2)
with sources.stream() as events_gen:
    for event in events_gen:
        ...  # Does event come from gen1 or gen2?

then I have to test the data to determine which of the source generators this event — I hope it's okay I refer to the data as events — was taken from. This is not always possible. For instance, if I have two sources generating strings, there is no sure way to discern which source stream generated a given string.

Would you consider something like this in addition?

sources = aiostream.stream.mergetuple(gen1, gen2)
with sources.stream() as events_gen:
    for gen, event in events_gen:
        if gen == gen1:
            # event comes from gen1
        elif gen == gen2:
            # event comes from gen2

Or is there another way?

I am currently using ziplatest and keeping track of each object so that I can identify which of the items in the zip actually changed vs. the last iteration. This is quite ugly, I find.

Looking forward to ideas, or a sign that you'd be open to a PR…

[vxgmichel]

Hi @madduck, and thanks for the report :)

I think this can already be done quite simply:

from asyncio import run

from aiostream import stream, pipe


async def main():
    xs = stream.iterate("ABCDEF") | pipe.spaceout(0.05)
    ys = stream.iterate("abce") | pipe.spaceout(0.08)
    zs = stream.merge(
        xs | pipe.map(lambda v: (0, v)),
        ys | pipe.map(lambda v: (1, v)),
    )

    async with zs.stream() as streamer:
        async for item in streamer:
            match item:
                case (0, v):
                    print(f"from xs: {v}")
                case (1, v):
                    print(f"from ys: {v}")


if __name__ == "__main__":
    run(main())

If you really want a dedicated operator, it can be written as such:

from asyncio import run
from typing import AsyncIterable, AsyncIterator
from aiostream import stream, pipable_operator, pipe


@pipable_operator
def enumerate_and_merge[T](
    x: AsyncIterable[T], *xs: AsyncIterable[T]
) -> AsyncIterator[tuple[int, T]]:
    return stream.merge.raw(
        *(stream.map.raw(x, lambda v, i=i: (i, v)) for i, x in enumerate((x, *xs)))
    )


async def main():
    xs = stream.iterate("ABCDEF") | pipe.spaceout(0.05)
    ys = stream.iterate("abcde") | pipe.spaceout(0.08)
    zs = enumerate_and_merge(xs, ys)
    async with zs.stream() as streamer:
        async for item in streamer:
            match item:
                case (0, v):
                    print(f"from xs: {v}")
                case (1, v):
                    print(f"from ys: {v}")


if __name__ == "__main__":
    run(main())