Make getting PhotoSize byte count more reusable internally

tejaswinilakshmi · Oct 31, 2019 · a67c947 · a67c947
1 parent 7e34618
commit a67c947
Show file tree

Hide file tree

Showing 2 changed files with 28 additions and 22 deletions.
diff --git a/telethon/tl/custom/file.py b/telethon/tl/custom/file.py
@@ -122,21  122,10 @@ def size(self):
         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
@@ -14,6  14,7 @@
 import os
 import re
 import struct
 from collections import namedtuple
 from mimetypes import guess_extension
 from types import GeneratorType
 
@@ -67,6  68,8 @@
     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