app.plugins.files: refactor as per core

app/plugins/files/download.py

@@ -2,8 +2,13 @@ import asyncio
 import os
 import time
 
-from ub_core.utils import (Download, DownloadedFile, bytes_to_mb,
-                           get_tg_media_details, progress)
+from ub_core.utils import (
+    Download,
+    DownloadedFile,
+    bytes_to_mb,
+    get_tg_media_details,
+    progress,
+)
 
 from app import BOT, Message, bot
 
@@ -31,6 +36,7 @@ async def down_load(bot: BOT, message: Message):
     await response.edit("Input verified....Starting Download...")
 
     file_name = None
+    dl_obj: None = None
 
     if message.replied and message.replied.media:
 
@@ -76,6 +82,9 @@ async def down_load(bot: BOT, message: Message):
 
     except Exception as e:
         await response.edit(str(e))
+    finally:
+        if dl_obj:
+            await dl_obj.close()
 
 
 async def telegram_download(
@@ -90,13 +99,9 @@ async def telegram_download(
     """
     tg_media = get_tg_media_details(message)
 
-    tg_media.file_name = file_name or tg_media.file_name
+    file_name = file_name or tg_media.file_name
 
-    media_obj: DownloadedFile = DownloadedFile(
-        name=tg_media.file_name,
-        dir=dir_name,
-        size=bytes_to_mb(tg_media.file_size),
-    )
+    media_obj: DownloadedFile = DownloadedFile(file=os.path.join(dir_name, file_name), size=tg_media.file_size)
 
     progress_args = (response, "Downloading...", media_obj.path)
 

app/plugins/files/rename.py

@@ -34,7 +34,7 @@ async def rename(bot: BOT, message: Message):
     await response.edit("Input verified....Starting Download...")
 
     if message.replied:
-
+        dl_obj: None = None
         download_coro = telegram_download(
             message=message.replied,
             dir_name=dl_path,
@@ -62,3 +62,7 @@ async def rename(bot: BOT, message: Message):
 
     except Exception as e:
         await response.edit(str(e))
+
+    finally:
+        if dl_obj:
+            await dl_obj.close()

app/plugins/files/upload.py

@@ -9,14 +9,13 @@ from ub_core.utils import (
     Download,
     DownloadedFile,
     MediaType,
-    bytes_to_mb,
     check_audio,
     get_duration,
     progress,
     take_ss,
 )
 
-from app import BOT, Config, Message, bot
+from app import BOT, Config, Message
 
 UPLOAD_TYPES = Union[BOT.send_audio, BOT.send_document, BOT.send_photo, BOT.send_video]
 
@@ -70,16 +69,16 @@ FILE_TYPE_MAP = {
 }
 
 
-def file_check(file: str) -> bool:
+def file_exists(file: str) -> bool:
     return os.path.isfile(file)
 
 
-def check_size(size: int | float) -> bool:
-    limit = 3999 if bot.me.is_premium else 1999
-    return size < limit
+def size_over_limit(size: int | float, client: BOT) -> bool:
+    limit = 3999 if client.me.is_premium else 1999
+    return size > limit
 
 
-@bot.add_cmd(cmd="upload")
+@BOT.add_cmd(cmd="upload")
 async def upload(bot: BOT, message: Message):
     """
     CMD: UPLOAD
@@ -108,38 +107,39 @@ async def upload(bot: BOT, message: Message):
         await response.delete()
         return
 
-    elif input.startswith("http") and not file_check(input):
-        dl_obj: Download = await Download.setup(
-            url=input,
-            dir=os.path.join("downloads", str(time.time())),
-            message_to_edit=response,
-        )
-
-        if not check_size(dl_obj.size):
-            await response.edit("<b>Aborted</b>, File size exceeds TG Limits!!!")
-            return
+    elif input.startswith("http") and not file_exists(input):
 
         try:
-            await response.edit("URL detected in input, Starting Download....")
-            file: DownloadedFile = await dl_obj.download()
+            async with Download(
+                url=input,
+                dir=os.path.join("downloads", str(time.time())),
+                message_to_edit=response,
+            ) as dl_obj:
+                if size_over_limit(dl_obj.size, client=bot):
+                    await response.edit(
+                        "<b>Aborted</b>, File size exceeds TG Limits!!!"
+                    )
+                    return
+
+                await response.edit("URL detected in input, Starting Download....")
+                file: DownloadedFile = await dl_obj.download()
+
         except asyncio.exceptions.CancelledError:
             await response.edit("Cancelled...")
             return
+
         except TimeoutError:
             await response.edit("Download Timeout...")
             return
+
         except Exception as e:
             await response.edit(str(e))
             return
 
-    elif file_check(input):
-        file = DownloadedFile(
-            name=input,
-            dir=os.path.dirname(input),
-            size=bytes_to_mb(os.path.getsize(input)),
-        )
+    elif file_exists(input):
+        file = DownloadedFile(file=input)
 
-        if not check_size(file.size):
+        if size_over_limit(file.size, client=bot):
             await response.edit("<b>Aborted</b>, File size exceeds TG Limits!!!")
             return
 
@@ -162,7 +162,7 @@ async def bulk_upload(message: Message, response: Message):
     else:
         path_regex = os.path.join(message.filtered_input, "*")
 
-    file_list = [f for f in glob.glob(path_regex) if file_check(f)]
+    file_list = [f for f in glob.glob(path_regex) if file_exists(f)]
 
     if not file_list:
         await response.edit("Invalid Folder path/regex or Folder Empty")
@@ -172,13 +172,9 @@ async def bulk_upload(message: Message, response: Message):
 
     for file in file_list:
 
-        file_info = DownloadedFile(
-            name=os.path.basename(file),
-            dir=os.path.dirname(file),
-            size=bytes_to_mb(os.path.getsize(file)),
-        )
+        file_info = DownloadedFile(file=file)
 
-        if not check_size(file_info.size):
+        if size_over_limit(file_info.size, client=message._client):
             await response.reply(
                 f"Skipping {file_info.name} due to size exceeding limit."
             )
@@ -214,6 +210,7 @@ async def upload_to_tg(file: DownloadedFile, message: Message, response: Message
             caption=file.name,
         )
         await response.delete()
+
     except asyncio.exceptions.CancelledError:
         await response.edit("Cancelled....")
         raise