如何在anyio中合并流？探寻多流交错迭代的现成方案

AnyIO本身没有直接提供合并多流并按数据到达顺序交错读取的API，但用**任务组（Task Group）+ 队列（Queue）**的组合就能实现简洁优雅的方案，比自定义CombinedReceiveStream子类更贴合AnyIO的异步设计范式。

核心实现思路 #

给每个待读取的流单独启动一个异步任务，任务读取到数据后立刻将数据（可附带流标识，方便做类似annotate-output的标注）放入队列；主线程从队列中取出数据，自然就能按数据到达的先后顺序交错处理。

代码示例 #

import anyio
from anyio.streams.buffered import BufferedReceiveStream

async def merge_streams(*streams, stream_names=None):
    # 给每个流分配标识，默认用stream_0、stream_1...命名
    stream_names = stream_names or [f"stream_{i}" for i in range(len(streams))]
    queue = anyio.Queue()

    async def read_single_stream(stream, stream_id):
        # 包装成BufferedReceiveStream，方便按块读取数据
        buffered_stream = BufferedReceiveStream(stream)
        try:
            async for chunk in buffered_stream:
                # 将流标识和数据块放入队列
                await queue.put((stream_id, chunk))
        finally:
            # 用None标记当前流读取结束
            await queue.put((stream_id, None))

    # 启动所有流的读取任务
    async with anyio.create_task_group() as task_group:
        for stream, name in zip(streams, stream_names):
            task_group.start_soon(read_single_stream, stream, name)

    # 从队列中取出数据，直到所有流都结束
    active_stream_count = len(streams)
    while active_stream_count > 0:
        stream_id, chunk = await queue.get()
        if chunk is None:
            active_stream_count -= 1
            continue
        # 按顺序返回带标识的数据块
        yield stream_id, chunk

# 用法示例
async def main():
    # 模拟两个不同的流（实际可以是子进程输出、网络流等）
    async def mock_stream(name, delay):
        for i in range(3):
            await anyio.sleep(delay)
            yield f"{name}: message {i}".encode()

    stream1 = mock_stream("PROC1", 0.5)
    stream2 = mock_stream("PROC2", 0.3)

    async for stream_id, chunk in merge_streams(stream1, stream2):
        # 类似annotate-output的标注输出
        print(f"[{stream_id}] {chunk.decode().strip()}")

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

优势说明 #

1. 无需自定义流子类，直接用AnyIO原生组件构建，代码更简洁易维护；
2. 天然支持流的动态结束，自动处理流的关闭逻辑；
3. 流标识可自定义，轻松实现annotate-output的标注效果。

内容的提问来源于stack exchange，提问作者HoverHell