WIP: Video support

classes/pyroclient.py

@@ -8,11 +8,13 @@ from shutil import rmtree
 from time import time
 from traceback import format_exc
 from typing import List, Tuple, Union
+import aiofiles
 
 import pyrogram
 from aiohttp import ClientSession
 from bson import ObjectId
 from dateutil.relativedelta import relativedelta
+from classes.enums.submission_types import SubmissionType
 from libbot import json_read, json_write
 from libbot.i18n import BotLocale
 from libbot.i18n.sync import _
@@ -43,10 +45,14 @@ from classes.exceptions import (
 )
 from classes.pyrocommand import PyroCommand
 from modules.api_client import (
-    BodyPhotoUploadAlbumsAlbumPhotosPost,
+    BodyPhotoUpload,
+    BodyVideoUpload,
     File,
+    Photo,
+    Video,
     client,
     photo_upload,
+    video_upload,
 )
 from modules.database import col_submitted
 from modules.http_client import http_session
@@ -398,7 +404,7 @@ class PyroClient(Client):
                 language_code=command_set.language_code,
             )
 
-    async def submit_photo(
+    async def submit_media(
         self, id: str
     ) -> Tuple[Union[Message, None], Union[str, None]]:
         db_entry = col_submitted.find_one({"_id": ObjectId(id)})
@@ -431,24 +437,47 @@ class PyroClient(Client):
                     db_entry["user"], db_entry["telegram"]["msg_id"]
                 )
 
-        with open(str(filepath), "rb") as fh:
-            photo_bytes = BytesIO(fh.read())
+        async with aiofiles.open(str(filepath), "rb") as fh:
+            media_bytes = BytesIO(await fh.read())
 
         try:
-            response = await photo_upload(
-                self.config["posting"]["api"]["album"],
-                client=client,
-                multipart_data=BodyPhotoUploadAlbumsAlbumPhotosPost(
-                    File(photo_bytes, filepath.name, "image/jpeg")
-                ),
-                ignore_duplicates=self.config["submission"]["allow_duplicates"],
-                compress=False,
-                caption="queue",
-            )
+            if db_entry["type"] == SubmissionType.PHOTO.value:
+                response = await photo_upload(
+                    self.config["posting"]["api"]["album"],
+                    client=client,
+                    multipart_data=BodyPhotoUpload(
+                        File(media_bytes, filepath.name, "image/jpeg")
+                    ),
+                    ignore_duplicates=self.config["submission"]["allow_duplicates"],
+                    compress=False,
+                    caption="queue",
+                )
+            elif db_entry["type"] == SubmissionType.VIDEO.value:
+                response = await video_upload(
+                    self.config["posting"]["api"]["album"],
+                    client=client,
+                    multipart_data=BodyVideoUpload(
+                        File(media_bytes, filepath.name, "video/*")
+                    ),
+                    caption="queue",
+                )
+            elif db_entry["type"] == SubmissionType.ANIMATION.value:
+                response = await video_upload(
+                    self.config["posting"]["api"]["album"],
+                    client=client,
+                    multipart_data=BodyVideoUpload(
+                        File(media_bytes, filepath.name, "video/*")
+                    ),
+                    caption="queue",
+                )
         except UnexpectedStatus:
             raise SubmissionUnsupportedError(str(filepath))
 
-        response_dict = loads(response.content.decode("utf-8"))
+        response_dict = (
+            {}
+            if not hasattr(response, "content")
+            else loads(response.content.decode("utf-8"))
+        )
 
         if "duplicates" in response_dict and len(response_dict["duplicates"]) > 0:
             duplicates = []
@@ -480,7 +509,10 @@ class PyroClient(Client):
         except (FileNotFoundError, NotADirectoryError):
             logger.error("Could not delete '%s' on submission accepted", filepath)
 
-        return submission, response.parsed.id
+        return (
+            submission,
+            response.id if not hasattr(response, "parsed") else response.parsed.id,
+        )
 
     async def ban_user(self, id: int) -> None:
         pass