diff --git a/app/plugins/files/download.py b/app/plugins/files/download.py
index f8cc035..0b6b09a 100644
--- a/app/plugins/files/download.py
+++ b/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)
diff --git a/app/plugins/files/rename.py b/app/plugins/files/rename.py
index ab45fad..c444afc 100644
--- a/app/plugins/files/rename.py
+++ b/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()
diff --git a/app/plugins/files/upload.py b/app/plugins/files/upload.py
index 8a2dba0..d2445c7 100644
--- a/app/plugins/files/upload.py
+++ b/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("Aborted, 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(
+ "Aborted, 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("Aborted, 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