Make getting PhotoSize byte count more reusable internally

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:

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):
+def _photo_size_byte_count(size):
-    if len(stripped) < 3 or stripped[0] != 1:
+    if isinstance(size, types.PhotoSize):
-        return len(stripped)
+        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