如何使用asyncio从async for循环中yield数据？

我明白你在实现异步批量数据生成器时遇到的困惑——异步迭代器和普通同步迭代器的工作逻辑确实有差异，尤其是在async for循环中返回数据的方式。你的代码框架已经有了雏形，我来帮你补全并解释关键要点。

完整的异步批量生成器实现（推荐方式） #

这种方式使用异步生成器实现，代码更简洁易读，符合Python异步编程的惯用写法：

import asyncio
from typing import List, AsyncGenerator

class AsyncSimpleIterator:
    def __init__(self, data: List[str], batch_size: int = 1):
        self.data = data
        self.batch_size = batch_size
        self.current_index = 0

    async def get_batch_data(self, doc_ids: List[int]) -> List[str]:
        # 模拟异步IO操作（比如调用API、读取数据库）
        await asyncio.sleep(0.1)  # 用sleep模拟延迟
        return [self.data[idx] for idx in doc_ids]

    def __aiter__(self) -> AsyncGenerator[List[str], None]:
        # 返回异步生成器，支持async for迭代
        return self._async_generator()

    async def _async_generator(self) -> AsyncGenerator[List[str], None]:
        while self.current_index < len(self.data):
            # 计算当前批次的索引范围，避免越界
            end_idx = min(self.current_index + self.batch_size, len(self.data))
            batch_ids = list(range(self.current_index, end_idx))
            
            # 异步获取批量数据
            batch_data = await self.get_batch_data(batch_ids)
            
            # 异步yield数据，供async for循环接收
            yield batch_data
            
            # 更新迭代位置
            self.current_index = end_idx

关键逻辑解释 #

1. 异步迭代协议：
   要让类支持async for，必须实现异步迭代协议。这里通过__aiter__返回一个异步生成器（由async def+yield定义），Python会自动为生成器实现__anext__方法，无需手动处理StopAsyncIteration异常。
2. 批量控制：
   用current_index跟踪当前迭代位置，每次计算批次的起始/结束索引，确保最后一批数据不会超出列表长度。
3. 异步数据获取：
   get_batch_data中的await语句是异步逻辑的核心，用于处理IO密集型任务（比如网络请求），避免阻塞主线程。

使用示例 #

async def main():
    test_data = ["文档1", "文档2", "文档3", "文档4", "文档5"]
    # 初始化迭代器，设置批次大小为2
    iterator = AsyncSimpleIterator(test_data, batch_size=2)
    
    # 用async for循环迭代获取批量数据
    async for batch in iterator:
        print(f"收到批量数据: {batch}")

asyncio.run(main())

运行结果：

收到批量数据: ['文档1', '文档2']
收到批量数据: ['文档3', '文档4']
收到批量数据: ['文档5']

另一种实现方式：手动实现__anext__ #

如果你需要更精细地控制迭代过程，可以直接实现异步迭代器的__anext__方法：

import asyncio
from typing import List

class AsyncSimpleIterator:
    def __init__(self, data: List[str], batch_size: int = 1):
        self.data = data
        self.batch_size = batch_size
        self.current_index = 0

    async def get_batch_data(self, doc_ids: List[int]) -> List[str]:
        await asyncio.sleep(0.1)
        return [self.data[idx] for idx in doc_ids]

    def __aiter__(self):
        # 返回自身作为迭代器
        return self

    async def __anext__(self):
        # 终止迭代的条件
        if self.current_index >= len(self.data):
            raise StopAsyncIteration
        
        end_idx = min(self.current_index + self.batch_size, len(self.data))
        batch_ids = list(range(self.current_index, end_idx))
        batch_data = await self.get_batch_data(batch_ids)
        
        self.current_index = end_idx
        return batch_data

这种方式需要手动抛出StopAsyncIteration来终止迭代，适合需要自定义迭代终止逻辑的场景。

内容的提问来源于stack exchange，提问作者my-master