Python中带超时迭代AsyncIterator流且不终止流的缓冲实现咨询

我正在处理一个AsyncIterator[str]类型的对象，它从网络获取消息并以字符串形式生成。想要为这个流创建一个包装器，对消息进行缓冲并定期输出。

原代码如下：

import asyncio
import time
from typing import AsyncIterator, Optional

async def buffer_stream(stream: AsyncIterator[str], buffer_time: Optional[float]) -> AsyncIterator[str]:
    """
    Buffer messages from the stream, and yields them at regular intervals.
    """
    last_sent_at = time.perf_counter()
    buffer = ''

    stop = False
    while not stop:
        time_to_send = False

        timeout = (
            max(buffer_time - (time.perf_counter() - last_sent_at), 0)
            if buffer_time else None
        )
        try:
            buffer += await asyncio.wait_for(
                stream.__anext__(),
                timeout=timeout
            )
        except asyncio.TimeoutError:
            time_to_send = True
        except StopAsyncIteration:
            time_to_send = True
            stop = True
        else:
            if time.perf_counter() - last_sent_at >= buffer_time:
                time_to_send = True

        if not buffer_time or time_to_send:
            if buffer:
                yield buffer
                buffer = ''
            last_sent_at = time.perf_counter()

原逻辑看似合理，但第一次超时后流会被中断提前退出，无法处理完整流。尝试用asyncio.shield()包装stream.__anext__()避免任务取消，却引发RuntimeError: anext(): asynchronous generator is already running——因为前一个anext()仍在执行时又调用了它。

正确实现方案 #

核心问题在于异步生成器不支持并发调用__anext__()，不能直接对其使用shield。正确做法是将流的读取放在后台任务中持续运行，同时通过队列传递消息，主循环负责定时输出缓冲内容。

import asyncio
import time
from typing import AsyncIterator, Optional, Queue

async def buffer_stream(stream: AsyncIterator[str], buffer_time: Optional[float]) -> AsyncIterator[str]:
    """
    Buffer messages from the stream, and yields them at regular intervals.
    """
    queue: Queue[str] = asyncio.Queue()
    buffer = ''
    last_sent_at = time.perf_counter()

    # 后台任务：持续读取流并将消息放入队列
    async def read_stream():
        try:
            async for msg in stream:
                await queue.put(msg)
        finally:
            # 放入结束标记，通知主循环流已耗尽
            await queue.put(None)

    # 启动后台读取任务
    read_task = asyncio.create_task(read_stream())

    try:
        while True:
            # 计算超时时间：距离下次需要输出的剩余时长
            timeout = None
            if buffer_time is not None:
                elapsed = time.perf_counter() - last_sent_at
                timeout = max(buffer_time - elapsed, 0)

            try:
                # 等待队列消息或超时触发
                msg = await asyncio.wait_for(queue.get(), timeout=timeout)
            except asyncio.TimeoutError:
                # 超时后输出当前缓冲内容
                if buffer:
                    yield buffer
                    buffer = ''
                last_sent_at = time.perf_counter()
                continue

            if msg is None:
                # 流已结束，输出剩余缓冲内容后退出
                if buffer:
                    yield buffer
                break

            # 累积消息到缓冲
            buffer += msg

            # 检查是否达到缓冲时间阈值，触发输出
            if buffer_time is not None and (time.perf_counter() - last_sent_at >= buffer_time):
                yield buffer
                buffer = ''
                last_sent_at = time.perf_counter()
    finally:
        # 确保后台任务被清理，避免资源泄漏
        read_task.cancel()
        try:
            await read_task
        except asyncio.CancelledError:
            pass

实现说明 #

1. 后台读取任务：用async for安全遍历异步流，将消息放入队列，流结束时放入None标记。彻底避免了并发调用__anext__()的问题。
2. 队列+超时机制：主循环通过asyncio.wait_for等待队列消息，超时则输出缓冲；收到消息后累积，同时检查是否达到缓冲时间阈值，达到则立即输出。
3. 资源清理：通过finally块确保后台任务被取消并等待完成，防止异步资源泄漏。

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