如何在 Python3 中将异步生成器 stream 转换为类文件 object？

Question

所以我制作了一个网络服务（基于 starlette），其端点接受二进制主体。 我想将这个二进制体提供给 fastavro。

Starlette 文档说，我可以使用request.stream()作为异步 stream 访问原始数据。

async for chunk in request.stream():
    # do something with chunk...

现在，我想将 stream 提供给 fastavro。 问题是， fastavro 阅读器需要一个类似文件的输入 stream：

with open('some-file.avro', 'rb') as fo:
    avro_reader = reader(fo)

我的问题是，是否有一种干净的方法可以将此异步 stream 转换为类文件？

我想我可以实现一个 object，它有一个 read() 方法，等待并返回 request.stream 返回的数据。 但是如果调用者传递一个大小，我需要有一个 memory 缓冲区，不是吗？ 可以基于 BufferedRWPair 的东西吗？

或者是先将整个 stream 存储到磁盘或 memory，然后再将其提供给 fastavro 的唯一方法？

提前致谢 ！

Answer 1

我最终使用了 SpooledTemporaryFile：

data_file = SpooledTemporaryFile(mode='w+b',
        max_size=MAX_RECEIVED_DATA_MEMORY_SIZE)
async for chunk in request.stream():
    data_file.write(chunk)
data_file.seek(0)
avro_reader = reader(data_file)

这不是我设想的理想解决方案（以某种方式直接在输入和输出之间传输数据），但仍然足够好......

Answer 2

我遇到了同样的问题，写了 compact class StreamingBody 。 它正是我需要的。

from typing import AsyncIterator
import asyncio


class AsyncGen:
    def __init__(self, block_count, block_size) -> None:
        self.bc = block_count
        self.bs = block_size

    def __aiter__(self):
        return self

    async def __anext__(self):

        if self.bc == 0:
            raise StopAsyncIteration()

        self.bc -= 1
        return b"A" * self.bs


class StreamingBody:

    _chunks: AsyncIterator[bytes]
    _backlog: bytes

    def __init__(self, chunks: AsyncIterator[bytes]):
        self._chunks = chunks
        self._backlog = b""

    async def _read_until_end(self):

        content = self._backlog
        self._backlog = b""

        while True:
            try:
                content += await self._chunks.__anext__()
            except StopAsyncIteration:
                break

        return content

    async def _read_chunk(self, size: int):

        content = self._backlog
        bytes_read = len(self._backlog)

        while bytes_read < size:

            try:
                chunk = await self._chunks.__anext__()
            except StopAsyncIteration:
                break

            content += chunk
            bytes_read += len(chunk)

        self._backlog = content[size:]
        content = content[:size]

        return content

    async def read(self, size: int = -1):
        if size > 0:
            return await self._read_chunk(size)
        elif size == -1:
            return await self._read_until_end()
        else:
            return b""

async def main():
    async_gen = AsyncGen(11, 3)
    body = StreamingBody(async_gen)

    res = await body.read(11)
    print(f"[{len(res)}]: {res}")

    res = await body.read()
    print(f"[{len(res)}]: {res}")

    res = await body.read()
    print(f"[{len(res)}]: {res}")


loop = asyncio.get_event_loop()
loop.run_until_complete(main())