| 18 | |
| 19 | |
| 20 | class ByteStreamer: |
| 21 | def __init__(self, client: Client): |
| 22 | """A custom class that holds the cache of a specific client and class functions. |
| 23 | attributes: |
| 24 | client: the client that the cache is for. |
| 25 | cached_file_ids: a dict of cached file IDs. |
| 26 | cached_file_properties: a dict of cached file properties. |
| 27 | |
| 28 | functions: |
| 29 | generate_file_properties: returns the properties for a media of a specific message contained in Tuple. |
| 30 | generate_media_session: returns the media session for the DC that contains the media file. |
| 31 | yield_file: yield a file from telegram servers for streaming. |
| 32 | |
| 33 | This is a modified version of the <https://github.com/eyaadh/megadlbot_oss/blob/master/mega/telegram/utils/custom_download.py> |
| 34 | Thanks to Eyaadh <https://github.com/eyaadh> |
| 35 | """ |
| 36 | self.clean_timer = 30 * 60 |
| 37 | self.client: Client = client |
| 38 | self.cached_file_ids: Dict[int, FileId] = {} |
| 39 | asyncio.create_task(self.clean_cache()) |
| 40 | |
| 41 | async def get_file_properties(self, id: int) -> FileId: |
| 42 | """ |
| 43 | Returns the properties of a media of a specific message in a FIleId class. |
| 44 | if the properties are cached, then it'll return the cached results. |
| 45 | or it'll generate the properties from the Message ID and cache them. |
| 46 | """ |
| 47 | if id not in self.cached_file_ids: |
| 48 | await self.generate_file_properties(id) |
| 49 | logging.debug(f"Cached file properties for message with ID {id}") |
| 50 | return self.cached_file_ids[id] |
| 51 | |
| 52 | async def generate_file_properties(self, id: int) -> FileId: |
| 53 | """ |
| 54 | Generates the properties of a media file on a specific message. |
| 55 | returns ths properties in a FIleId class. |
| 56 | """ |
| 57 | file_id = await get_file_ids(self.client, LOG_CHANNEL, id) |
| 58 | logging.debug(f"Generated file ID and Unique ID for message with ID {id}") |
| 59 | if not file_id: |
| 60 | logging.debug(f"Message with ID {id} not found") |
| 61 | raise FIleNotFound |
| 62 | self.cached_file_ids[id] = file_id |
| 63 | logging.debug(f"Cached media message with ID {id}") |
| 64 | return self.cached_file_ids[id] |
| 65 | |
| 66 | async def generate_media_session(self, client: Client, file_id: FileId) -> Session: |
| 67 | """ |
| 68 | Generates the media session for the DC that contains the media file. |
| 69 | This is required for getting the bytes from Telegram servers. |
| 70 | """ |
| 71 | |
| 72 | media_session = client.media_sessions.get(file_id.dc_id, None) |
| 73 | |
| 74 | if media_session is None: |
| 75 | if file_id.dc_id != await client.storage.dc_id(): |
| 76 | media_session = Session( |
| 77 | client, |