Add option for parallel streamed file transfer

mautrix_telegram/util/file_transfer.py

@@ -33,6 +33,7 @@ from mautrix.appservice import IntentAPI
 from ..tgclient import MautrixTelegramClient
 from ..db import TelegramFile as DBTelegramFile
 from ..util import sane_mimetypes
+from .parallel_file_transfer import parallel_transfer_to_matrix
 
 try:
     from PIL import Image
@@ -126,7 +127,7 @@ async def transfer_thumbnail_to_matrix(client: MautrixTelegramClient, intent: In
         return db_file
 
     video_ext = sane_mimetypes.guess_extension(mime)
-    if VideoFileClip and video_ext:
+    if VideoFileClip and video_ext and video:
         try:
             file, width, height = _read_video_thumbnail(video, video_ext, frame_ext="png")
         except OSError:
@@ -158,7 +159,8 @@ TypeThumbnail = Optional[Union[TypeLocation, TypePhotoSize]]
 
 async def transfer_file_to_matrix(client: MautrixTelegramClient, intent: IntentAPI,
                                   location: TypeLocation, thumbnail: TypeThumbnail = None,
-                                  is_sticker: bool = False) -> Optional[DBTelegramFile]:
+                                  is_sticker: bool = False, filename: Optional[str] = None,
+                                  parallel_id: Optional[int] = None) -> Optional[DBTelegramFile]:
     location_id = _location_to_id(location)
     if not location_id:
         return None
@@ -174,43 +176,52 @@ async def transfer_file_to_matrix(client: MautrixTelegramClient, intent: IntentA
         transfer_locks[location_id] = lock
     async with lock:
         return await _unlocked_transfer_file_to_matrix(client, intent, location_id, location,
-                                                       thumbnail, is_sticker)
+                                                       thumbnail, is_sticker, filename,
+                                                       parallel_id)
 
 
 async def _unlocked_transfer_file_to_matrix(client: MautrixTelegramClient, intent: IntentAPI,
                                             loc_id: str, location: TypeLocation,
-                                            thumbnail: TypeThumbnail, is_sticker: bool
+                                            thumbnail: TypeThumbnail, is_sticker: bool,
+                                            filename: Optional[str],
+                                            parallel_id: Optional[int] = None
                                             ) -> Optional[DBTelegramFile]:
     db_file = DBTelegramFile.get(loc_id)
     if db_file:
         return db_file
 
-    try:
-        file = await client.download_file(location)
-    except (LocationInvalidError, FileIdInvalidError):
-        return None
-    except (AuthBytesInvalidError, AuthKeyInvalidError, SecurityError) as e:
-        log.exception(f"{e.__class__.__name__} while downloading a file.")
-        return None
+    if parallel_id and isinstance(location, Document):
+        db_file = await parallel_transfer_to_matrix(client, intent, loc_id, location, filename,
+                                                    parallel_id)
+        mime_type = location.mime_type
+        file = None
+    else:
+        try:
+            file = await client.download_file(location)
+        except (LocationInvalidError, FileIdInvalidError):
+            return None
+        except (AuthBytesInvalidError, AuthKeyInvalidError, SecurityError) as e:
+            log.exception(f"{e.__class__.__name__} while downloading a file.")
+            return None
 
-    width, height = None, None
-    mime_type = magic.from_buffer(file, mime=True)
+        width, height = None, None
+        mime_type = magic.from_buffer(file, mime=True)
 
-    image_converted = False
-    if mime_type == "image/webp":
-        new_mime_type, file, width, height = convert_image(
-            file, source_mime="image/webp", target_type="png",
-            thumbnail_to=(256, 256) if is_sticker else None)
-        image_converted = new_mime_type != mime_type
-        mime_type = new_mime_type
-        thumbnail = None
+        image_converted = False
+        if mime_type == "image/webp":
+            new_mime_type, file, width, height = convert_image(
+                file, source_mime="image/webp", target_type="png",
+                thumbnail_to=(256, 256) if is_sticker else None)
+            image_converted = new_mime_type != mime_type
+            mime_type = new_mime_type
+            thumbnail = None
 
-    content_uri = await intent.upload_media(file, mime_type)
+        content_uri = await intent.upload_media(file, mime_type)
 
-    db_file = DBTelegramFile(id=loc_id, mxc=content_uri,
-                             mime_type=mime_type, was_converted=image_converted,
-                             timestamp=int(time.time()), size=len(file),
-                             width=width, height=height)
+        db_file = DBTelegramFile(id=loc_id, mxc=content_uri,
+                                 mime_type=mime_type, was_converted=image_converted,
+                                 timestamp=int(time.time()), size=len(file),
+                                 width=width, height=height)
     if thumbnail and (mime_type.startswith("video/") or mime_type == "image/gif"):
         if isinstance(thumbnail, (PhotoSize, PhotoCachedSize)):
             thumbnail = thumbnail.location