Spaces:
Running
Running
| # copied from https://github.com/tulir/mautrix-telegram/blob/master/mautrix_telegram/util/parallel_file_transfer.py | |
| # Copyright (C) 2021 Tulir Asokan | |
| import asyncio | |
| import hashlib | |
| import inspect | |
| import logging | |
| import math | |
| import os | |
| from collections import defaultdict | |
| from typing import Optional, List, AsyncGenerator, Union, Awaitable, DefaultDict, Tuple, BinaryIO | |
| from telethon import utils, helpers, TelegramClient | |
| from telethon.crypto import AuthKey | |
| from telethon.network import MTProtoSender | |
| from telethon.tl.alltlobjects import LAYER | |
| from telethon.tl.functions import InvokeWithLayerRequest | |
| from telethon.tl.functions.auth import ExportAuthorizationRequest, ImportAuthorizationRequest | |
| from telethon.tl.functions.upload import (GetFileRequest, SaveFilePartRequest, | |
| SaveBigFilePartRequest) | |
| from telethon.tl.types import (Document, InputFileLocation, InputDocumentFileLocation, | |
| InputPhotoFileLocation, InputPeerPhotoFileLocation, TypeInputFile, | |
| InputFileBig, InputFile) | |
| try: | |
| from mautrix.crypto.attachments import async_encrypt_attachment | |
| except ImportError: | |
| async_encrypt_attachment = None | |
| log: logging.Logger = logging.getLogger("telethon") | |
| TypeLocation = Union[Document, InputDocumentFileLocation, InputPeerPhotoFileLocation, | |
| InputFileLocation, InputPhotoFileLocation] | |
| class DownloadSender: | |
| client: TelegramClient | |
| sender: MTProtoSender | |
| request: GetFileRequest | |
| remaining: int | |
| stride: int | |
| def __init__(self, client: TelegramClient, sender: MTProtoSender, file: TypeLocation, offset: int, limit: int, | |
| stride: int, count: int) -> None: | |
| self.sender = sender | |
| self.client = client | |
| self.request = GetFileRequest(file, offset=offset, limit=limit) | |
| self.stride = stride | |
| self.remaining = count | |
| async def next(self) -> Optional[bytes]: | |
| if not self.remaining: | |
| return None | |
| result = await self.client._call(self.sender, self.request) | |
| self.remaining -= 1 | |
| self.request.offset += self.stride | |
| return result.bytes | |
| def disconnect(self) -> Awaitable[None]: | |
| return self.sender.disconnect() | |
| class UploadSender: | |
| client: TelegramClient | |
| sender: MTProtoSender | |
| request: Union[SaveFilePartRequest, SaveBigFilePartRequest] | |
| part_count: int | |
| stride: int | |
| previous: Optional[asyncio.Task] | |
| loop: asyncio.AbstractEventLoop | |
| def __init__(self, client: TelegramClient, sender: MTProtoSender, file_id: int, part_count: int, big: bool, | |
| index: int, | |
| stride: int, loop: asyncio.AbstractEventLoop) -> None: | |
| self.client = client | |
| self.sender = sender | |
| self.part_count = part_count | |
| if big: | |
| self.request = SaveBigFilePartRequest(file_id, index, part_count, b"") | |
| else: | |
| self.request = SaveFilePartRequest(file_id, index, b"") | |
| self.stride = stride | |
| self.previous = None | |
| self.loop = loop | |
| async def next(self, data: bytes) -> None: | |
| if self.previous: | |
| await self.previous | |
| self.previous = self.loop.create_task(self._next(data)) | |
| async def _next(self, data: bytes) -> None: | |
| self.request.bytes = data | |
| log.debug(f"Sending file part {self.request.file_part}/{self.part_count}" | |
| f" with {len(data)} bytes") | |
| await self.client._call(self.sender, self.request) | |
| self.request.file_part += self.stride | |
| async def disconnect(self) -> None: | |
| if self.previous: | |
| await self.previous | |
| return await self.sender.disconnect() | |
| class ParallelTransferrer: | |
| client: TelegramClient | |
| loop: asyncio.AbstractEventLoop | |
| dc_id: int | |
| senders: Optional[List[Union[DownloadSender, UploadSender]]] | |
| auth_key: AuthKey | |
| upload_ticker: int | |
| def __init__(self, client: TelegramClient, dc_id: Optional[int] = None) -> None: | |
| self.client = client | |
| self.loop = self.client.loop | |
| self.dc_id = dc_id or self.client.session.dc_id | |
| self.auth_key = (None if dc_id and self.client.session.dc_id != dc_id | |
| else self.client.session.auth_key) | |
| self.senders = None | |
| self.upload_ticker = 0 | |
| async def _cleanup(self) -> None: | |
| await asyncio.gather(*[sender.disconnect() for sender in self.senders]) | |
| self.senders = None | |
| def _get_connection_count(file_size: int, max_count: int = 20, | |
| full_size: int = 100 * 1024 * 1024) -> int: | |
| if file_size > full_size: | |
| return max_count | |
| return math.ceil((file_size / full_size) * max_count) | |
| async def _init_download(self, connections: int, file: TypeLocation, part_count: int, | |
| part_size: int) -> None: | |
| minimum, remainder = divmod(part_count, connections) | |
| def get_part_count() -> int: | |
| nonlocal remainder | |
| if remainder > 0: | |
| remainder -= 1 | |
| return minimum + 1 | |
| return minimum | |
| # The first cross-DC sender will export+import the authorization, so we always create it | |
| # before creating any other senders. | |
| self.senders = [ | |
| await self._create_download_sender(file, 0, part_size, connections * part_size, | |
| get_part_count()), | |
| *await asyncio.gather( | |
| *[self._create_download_sender(file, i, part_size, connections * part_size, | |
| get_part_count()) | |
| for i in range(1, connections)]) | |
| ] | |
| async def _create_download_sender(self, file: TypeLocation, index: int, part_size: int, | |
| stride: int, | |
| part_count: int) -> DownloadSender: | |
| return DownloadSender(self.client, await self._create_sender(), file, index * part_size, part_size, | |
| stride, part_count) | |
| async def _init_upload(self, connections: int, file_id: int, part_count: int, big: bool | |
| ) -> None: | |
| self.senders = [ | |
| await self._create_upload_sender(file_id, part_count, big, 0, connections), | |
| *await asyncio.gather( | |
| *[self._create_upload_sender(file_id, part_count, big, i, connections) | |
| for i in range(1, connections)]) | |
| ] | |
| async def _create_upload_sender(self, file_id: int, part_count: int, big: bool, index: int, | |
| stride: int) -> UploadSender: | |
| return UploadSender(self.client, await self._create_sender(), file_id, part_count, big, index, stride, | |
| loop=self.loop) | |
| async def _create_sender(self) -> MTProtoSender: | |
| dc = await self.client._get_dc(self.dc_id) | |
| sender = MTProtoSender(self.auth_key, loggers=self.client._log) | |
| await sender.connect(self.client._connection(dc.ip_address, dc.port, dc.id, | |
| loggers=self.client._log, | |
| proxy=self.client._proxy)) | |
| if not self.auth_key: | |
| log.debug(f"Exporting auth to DC {self.dc_id}") | |
| auth = await self.client(ExportAuthorizationRequest(self.dc_id)) | |
| self.client._init_request.query = ImportAuthorizationRequest(id=auth.id, | |
| bytes=auth.bytes) | |
| req = InvokeWithLayerRequest(LAYER, self.client._init_request) | |
| await sender.send(req) | |
| self.auth_key = sender.auth_key | |
| return sender | |
| async def init_upload(self, file_id: int, file_size: int, part_size_kb: Optional[float] = None, | |
| connection_count: Optional[int] = None) -> Tuple[int, int, bool]: | |
| connection_count = connection_count or self._get_connection_count(file_size) | |
| part_size = (part_size_kb or utils.get_appropriated_part_size(file_size)) * 1024 | |
| part_count = (file_size + part_size - 1) // part_size | |
| is_large = file_size > 10 * 1024 * 1024 | |
| await self._init_upload(connection_count, file_id, part_count, is_large) | |
| return part_size, part_count, is_large | |
| async def upload(self, part: bytes) -> None: | |
| await self.senders[self.upload_ticker].next(part) | |
| self.upload_ticker = (self.upload_ticker + 1) % len(self.senders) | |
| async def finish_upload(self) -> None: | |
| await self._cleanup() | |
| async def download(self, file: TypeLocation, file_size: int, | |
| part_size_kb: Optional[float] = None, | |
| connection_count: Optional[int] = None) -> AsyncGenerator[bytes, None]: | |
| connection_count = connection_count or self._get_connection_count(file_size) | |
| part_size = (part_size_kb or utils.get_appropriated_part_size(file_size)) * 1024 | |
| part_count = math.ceil(file_size / part_size) | |
| log.debug("Starting parallel download: " | |
| f"{connection_count} {part_size} {part_count} {file!s}") | |
| await self._init_download(connection_count, file, part_count, part_size) | |
| part = 0 | |
| while part < part_count: | |
| tasks = [] | |
| for sender in self.senders: | |
| tasks.append(self.loop.create_task(sender.next())) | |
| for task in tasks: | |
| data = await task | |
| if not data: | |
| break | |
| yield data | |
| part += 1 | |
| log.debug(f"Part {part} downloaded") | |
| log.debug("Parallel download finished, cleaning up connections") | |
| await self._cleanup() | |
| parallel_transfer_locks: DefaultDict[int, asyncio.Lock] = defaultdict(lambda: asyncio.Lock()) | |
| def stream_file(file_to_stream: BinaryIO, chunk_size=1024): | |
| while True: | |
| data_read = file_to_stream.read(chunk_size) | |
| if not data_read: | |
| break | |
| yield data_read | |
| async def _internal_transfer_to_telegram(client: TelegramClient, | |
| response: BinaryIO, | |
| progress_callback: callable | |
| ) -> Tuple[TypeInputFile, int]: | |
| file_id = helpers.generate_random_long() | |
| file_size = os.path.getsize(response.name) | |
| hash_md5 = hashlib.md5() | |
| uploader = ParallelTransferrer(client) | |
| part_size, part_count, is_large = await uploader.init_upload(file_id, file_size) | |
| buffer = bytearray() | |
| for data in stream_file(response): | |
| if progress_callback: | |
| r = progress_callback(response.tell(), file_size) | |
| if inspect.isawaitable(r): | |
| await r | |
| if not is_large: | |
| hash_md5.update(data) | |
| if len(buffer) == 0 and len(data) == part_size: | |
| await uploader.upload(data) | |
| continue | |
| new_len = len(buffer) + len(data) | |
| if new_len >= part_size: | |
| cutoff = part_size - len(buffer) | |
| buffer.extend(data[:cutoff]) | |
| await uploader.upload(bytes(buffer)) | |
| buffer.clear() | |
| buffer.extend(data[cutoff:]) | |
| else: | |
| buffer.extend(data) | |
| if len(buffer) > 0: | |
| await uploader.upload(bytes(buffer)) | |
| await uploader.finish_upload() | |
| if is_large: | |
| return InputFileBig(file_id, part_count, "upload"), file_size | |
| else: | |
| return InputFile(file_id, part_count, "upload", hash_md5.hexdigest()), file_size | |
| async def download_file(client: TelegramClient, | |
| location: TypeLocation, | |
| out: BinaryIO, | |
| progress_callback: callable = None | |
| ) -> BinaryIO: | |
| size = location.size | |
| dc_id, location = utils.get_input_location(location) | |
| # We lock the transfers because telegram has connection count limits | |
| downloader = ParallelTransferrer(client, dc_id) | |
| downloaded = downloader.download(location, size) | |
| async for x in downloaded: | |
| out.write(x) | |
| if progress_callback: | |
| r = progress_callback(out.tell(), size) | |
| if inspect.isawaitable(r): | |
| await r | |
| return out | |
| async def upload_file(client: TelegramClient, | |
| file: BinaryIO, | |
| progress_callback: callable = None, | |
| ) -> TypeInputFile: | |
| res = (await _internal_transfer_to_telegram(client, file, progress_callback))[0] | |
| return res | |