From a67c94787b9c22ace4a9c704217d3e2a00ad9a7d Mon Sep 17 00:00:00 2001 From: Lonami Exo Date: Thu, 31 Oct 2019 19:38:27 +0100 Subject: [PATCH] Make getting PhotoSize byte count more reusable internally --- telethon/tl/custom/file.py | 13 +------------ telethon/utils.py | 37 +++++++++++++++++++++++++++---------- 2 files changed, 28 insertions(+), 22 deletions(-) diff --git a/telethon/tl/custom/file.py b/telethon/tl/custom/file.py index 69e9ef50..d1e5d834 100644 --- a/telethon/tl/custom/file.py +++ b/telethon/tl/custom/file.py @@ -122,21 +122,10 @@ class File: The size in bytes of this file. """ if isinstance(self.media, types.Photo): - return self._size_for(self.media.sizes[-1]) + return utils._photo_size_byte_count(self.media.sizes[-1]) elif isinstance(self.media, types.Document): return self.media.size - @staticmethod - def _size_for(kind): - if isinstance(kind, types.PhotoSize): - return kind.size - elif isinstance(kind, types.PhotoStrippedSize): - return utils._stripped_real_length(kind.bytes) - elif isinstance(kind, types.PhotoCachedSize): - return len(kind.bytes) - # elif isinstance(kind, types.PhotoSizeEmpty): - return 0 - def _from_attr(self, cls, field): if isinstance(self.media, types.Document): for attr in self.media.attributes: diff --git a/telethon/utils.py b/telethon/utils.py index 0505a215..40ab9e3e 100644 --- a/telethon/utils.py +++ b/telethon/utils.py @@ -14,6 +14,7 @@ import mimetypes import os import re import struct +from collections import namedtuple from mimetypes import guess_extension from types import GeneratorType @@ -67,6 +68,8 @@ VALID_USERNAME_RE = re.compile( re.IGNORECASE ) +_FileInfo = namedtuple('FileInfo', 'dc_id location size') + _log = logging.getLogger(__name__) @@ -676,9 +679,14 @@ def get_input_location(location): Note that this returns a tuple ``(dc_id, location)``, the ``dc_id`` being present if known. """ + info = _get_file_info(location) + return info.dc_id, info.location + + +def _get_file_info(location): try: if location.SUBCLASS_OF_ID == 0x1523d462: - return None, location # crc32(b'InputFileLocation'): + return _FileInfo(None, location, None) # crc32(b'InputFileLocation'): except AttributeError: _raise_cast_fail(location, 'InputFileLocation') @@ -691,19 +699,19 @@ def get_input_location(location): location = location.photo if isinstance(location, types.Document): - return (location.dc_id, types.InputDocumentFileLocation( + return _FileInfo(location.dc_id, types.InputDocumentFileLocation( id=location.id, access_hash=location.access_hash, file_reference=location.file_reference, thumb_size='' # Presumably to download one of its thumbnails - )) + ), location.size) elif isinstance(location, types.Photo): - return (location.dc_id, types.InputPhotoFileLocation( + return _FileInfo(location.dc_id, types.InputPhotoFileLocation( id=location.id, access_hash=location.access_hash, file_reference=location.file_reference, thumb_size=location.sizes[-1].type - )) + ), _photo_size_byte_count(location.sizes[-1])) if isinstance(location, types.FileLocationToBeDeprecated): raise TypeError('Unavailable location cannot be used as input') @@ -1298,7 +1306,7 @@ def stripped_photo_to_jpg(stripped): Ported from https://github.com/telegramdesktop/tdesktop/blob/bec39d89e19670eb436dc794a8f20b657cb87c71/Telegram/SourceFiles/ui/image/image.cpp#L225 """ - # NOTE: Changes here should update _stripped_real_length + # NOTE: Changes here should update _photo_size_byte_count if len(stripped) < 3 or stripped[0] != 1: return stripped @@ -1309,8 +1317,17 @@ def stripped_photo_to_jpg(stripped): return bytes(header) + stripped[3:] + footer -def _stripped_real_length(stripped): - if len(stripped) < 3 or stripped[0] != 1: - return len(stripped) +def _photo_size_byte_count(size): + if isinstance(size, types.PhotoSize): + return size.size + elif isinstance(size, types.PhotoStrippedSize): + if len(size.bytes) < 3 or size.bytes[0] != 1: + return len(size.bytes) - return len(stripped) + 622 + return len(size.bytes) + 622 + elif isinstance(size, types.PhotoCachedSize): + return len(size.bytes) + elif isinstance(size, types.PhotoSizeEmpty): + return 0 + else: + return None