tranny / App /Streaming /Utils /Download.py
Mbonea's picture
added streaming
07fd3f6
raw
history blame
2.78 kB
import math, asyncio, subprocess
from telethon import TelegramClient
from fastapi.responses import StreamingResponse
import logging
logging.getLogger(__name__)
logging.basicConfig(level=logging.INFO)
class Download:
client: TelegramClient
route: str
offset: int
handler: None
file: None
limit: int
file_size: float
def __init__(self, handler):
self.handler = handler
self.file = handler.message.media
self.file_size = handler.message.file.size
self.limit = handler.sanity.limit
self.offset = handler.sanity.offset
self.client = handler.client
self.mime_type = handler.message.file.mime_type
async def download(self):
part_size = int(512 * 1024) * 2
first_part_cut = self.offset % part_size
first_part = math.floor(self.offset / part_size)
last_part_cut = part_size - (self.limit % part_size)
last_part = math.ceil(self.limit / part_size)
part_count = math.ceil(self.file_size / part_size)
part = first_part
try:
async for chunk in self.client.iter_download(
self.file, offset=first_part * part_size, request_size=part_size
):
if part == first_part:
yield bytes(chunk[first_part_cut:])
elif part == last_part:
yield bytes(chunk[:last_part_cut])
else:
yield bytes(chunk)
logging.debug(f"Part {part}/{last_part} (total {part_count}) served!")
part += 1
logging.debug("serving finished")
except (GeneratorExit, StopAsyncIteration, asyncio.CancelledError):
logging.debug("file serve interrupted")
raise
except Exception as e:
print(e)
logging.debug("file serve errored", exc_info=True)
async def handle_request(self):
headers = {
"content-type": self.mime_type,
"content-range": f"bytes {self.offset}-{self.limit-1}/{self.file_size}",
"content-length": str(self.limit - self.offset),
"accept-ranges": "bytes",
"content-transfer-encoding": "Binary",
"content-disposition": f'{self.handler.route}; filename="{self.handler.message.file.name}"',
}
logging.info(
f"Serving file in {self.handler.message.file.name}) ; Range: {self.offset} - {self.limit}"
)
if self.handler.head:
body = None
else:
body = self.download()
return StreamingResponse(
media_type=self.mime_type,
content=body,
headers=headers,
status_code=206 if self.offset else 200,
)