Stream and reupload a file with pyrogram on the fly

pyrogram_stream.py

import os
from typing import BinaryIO

from pyrogram import Client, filters
from pyrogram.types import Message

bot = Client(
    "stream",
    api_id=int(os.getenv("API_ID")),
    api_hash=os.getenv("API_HASH"),
    bot_token=os.getenv("BOT_TOKEN"),
)


class Stream(io.BytesIO):
    def __init__(self, name: str, file_size: int, stream: BinaryIO, chunk_size: int = 512 * 1024):
        super().__init__()
        self.buffer = b""
        self.stream = stream.__aiter__()
        self.chunk_size = chunk_size
        self.name = name
        self.file_size = file_size
    
    def read(self, n):
        try:
            return self.buffer[:n]
        finally:
            self.buffer = self.buffer[n:]

    async def progress(self, current, total):
        if len(self.buffer) < self.chunk_size and current + self.chunk_size <= total:
            await self.fill()

    async def fill(self):
        self.buffer += await self.stream.__anext__()

    def tell(self) -> int:
        return self.file_size

    def seek(self, n, seek_type=None):
        pass


@bot.on_message(filters.private & filters.document)
async def reupload(bot: Client, message: Message):
    stream = Stream(
        name=message.document.file_name,
        file_size=message.document.file_size,
        stream=bot.stream_media(message)
    )
    await stream.fill()
    
    await message.reply_document(
        document=stream,
        caption="Here's your streamed file!",
        progress=stream.progress,
    )


bot.run()