refactor media group uploads

Author: nlef

bot/notifications.py

@@ -8,8 +8,7 @@
 from io import BytesIO
 import logging
 import re
-from typing import TYPE_CHECKING, Callable
-from typing_extensions import Final
+from typing import TYPE_CHECKING, Callable, Final
 
 import aiofiles
 import anyio
@@ -489,9 +488,16 @@ def _parse_path(ws_message: str) -> list[str]:
             path = [""]
         return path
 
-    async def _send_image(self, paths: list[str], message: str) -> None:
+    async def _send_media(
+        self,
+        paths: list[str],
+        message: str,
+        media_class: type[InputMediaPhoto | InputMediaVideo | InputMediaDocument],
+        media_name: str,
+        error_prefix: str,
+    ) -> None:
         try:
-            photos_list: list[InputMediaAudio | InputMediaDocument | InputMediaPhoto | InputMediaVideo] = []
+            media_list: list[InputMediaAudio | InputMediaDocument | InputMediaPhoto | InputMediaVideo] = []
             for path in paths:
                 path_obj = anyio.Path(path)
                 if not await path_obj.is_file():
@@ -505,96 +511,39 @@ async def _send_image(self, paths: list[str], message: str) -> None:
                     bio.write(await fh.read())
                 bio.seek(0)
                 if bio.getbuffer().nbytes > self._max_upload_file_size * 1024 * 1024:
-                    await self._bot.send_message(self._chat_id, text=f"Telegram bots have a {self._max_upload_file_size}mb filesize restriction, image couldn't be uploaded: `{path}`")
-                elif not photos_list:
-                    photos_list.append(InputMediaPhoto(bio, filename=bio.name, caption=message))
+                    await self._bot.send_message(self._chat_id, text=f"Telegram bots have a {self._max_upload_file_size}mb filesize restriction, {media_name} couldn't be uploaded: `{path}`")
+                elif not media_list:
+                    media_list.append(media_class(bio, filename=bio.name, caption=message))
                 else:
-                    photos_list.append(InputMediaPhoto(bio, filename=bio.name))
+                    media_list.append(media_class(bio, filename=bio.name))
                 bio.close()
 
             await self._bot.send_media_group(
                 self._chat_id,
-                media=photos_list,
+                media=media_list,
                 disable_notification=self._silent_commands,
+                write_timeout=120,
             )
 
         except Exception as ex:
             logger.warning(ex)
-            await self._bot.send_message(self._chat_id, text=f"Error sending image: {ex}", disable_notification=self._silent_commands)
+            await self._bot.send_message(self._chat_id, text=f"Error sending {error_prefix}: {ex}", disable_notification=self._silent_commands)
 
-    def send_image(self, ws_message: str) -> None:
-        self._schedule_job(self._send_image, {"paths": self._parse_path(ws_message), "message": self._parse_message(ws_message)})
+    async def _send_image(self, paths: list[str], message: str) -> None:
+        await self._send_media(paths, message, InputMediaPhoto, "image", "image")
 
     async def _send_video(self, paths: list[str], message: str) -> None:
-        try:
-            photos_list: list[InputMediaAudio | InputMediaDocument | InputMediaPhoto | InputMediaVideo] = []
-            for path in paths:
-                path_obj = anyio.Path(path)
-                if not await path_obj.is_file():
-                    await self._bot.send_message(self._chat_id, text="Provided path is not a file", disable_notification=self._silent_commands)
-                    return
-
-                bio = BytesIO()
-                bio.name = path_obj.name
-
-                async with aiofiles.open(path_obj, "rb") as fh:
-                    bio.write(await fh.read())
-                bio.seek(0)
-                if bio.getbuffer().nbytes > self._max_upload_file_size * 1024 * 1024:
-                    await self._bot.send_message(self._chat_id, text=f"Telegram bots have a {self._max_upload_file_size}mb filesize restriction, video couldn't be uploaded: `{path}`")
-                elif not photos_list:
-                    photos_list.append(InputMediaVideo(bio, filename=bio.name, caption=message))
-                else:
-                    photos_list.append(InputMediaVideo(bio, filename=bio.name))
-                bio.close()
+        await self._send_media(paths, message, InputMediaVideo, "video", "video")
 
-            await self._bot.send_media_group(
-                self._chat_id,
-                media=photos_list,
-                disable_notification=self._silent_commands,
-                write_timeout=120,
-            )
+    async def _send_document(self, paths: list[str], message: str) -> None:
+        await self._send_media(paths, message, InputMediaDocument, "document", "document")
 
-        except Exception as ex:
-            logger.warning(ex)
-            await self._bot.send_message(self._chat_id, text=f"Error sending video: {ex}", disable_notification=self._silent_commands)
+    def send_image(self, ws_message: str) -> None:
+        self._schedule_job(self._send_image, {"paths": self._parse_path(ws_message), "message": self._parse_message(ws_message)})
 
     def send_video(self, ws_message: str) -> None:
         self._schedule_job(self._send_video, {"paths": self._parse_path(ws_message), "message": self._parse_message(ws_message)})
 
-    async def _send_document(self, paths: list[str], message: str) -> None:
-        try:
-            photos_list: list[InputMediaAudio | InputMediaDocument | InputMediaPhoto | InputMediaVideo] = []
-            for path in paths:
-                path_obj = anyio.Path(path)
-                if not await path_obj.is_file():
-                    await self._bot.send_message(self._chat_id, text="Provided path is not a file", disable_notification=self._silent_commands)
-                    return
-
-                bio = BytesIO()
-                bio.name = path_obj.name
-
-                async with aiofiles.open(path_obj, "rb") as fh:
-                    bio.write(await fh.read())
-                bio.seek(0)
-                if bio.getbuffer().nbytes > self._max_upload_file_size * 1024 * 1024:
-                    await self._bot.send_message(self._chat_id, text=f"Telegram bots have a {self._max_upload_file_size}mb filesize restriction, document couldn't be uploaded: `{path}`")
-                elif not photos_list:
-                    photos_list.append(InputMediaDocument(bio, filename=bio.name, caption=message))
-                else:
-                    photos_list.append(InputMediaDocument(bio, filename=bio.name))
-                bio.close()
-
-            await self._bot.send_media_group(
-                self._chat_id,
-                media=photos_list,
-                disable_notification=self._silent_commands,
-            )
-
-        except Exception as ex:
-            logger.warning(ex)
-            await self._bot.send_message(self._chat_id, text=f"Error sending document: {ex}", disable_notification=self._silent_commands)
-
     def send_document(self, ws_message: str) -> None:
         self._schedule_job(self._send_document, {"paths": self._parse_path(ws_message), "message": self._parse_message(ws_message)})