diff --git a/src/engine/base.py b/src/engine/base.py index 906c6d1c..14ccd0d1 100644 --- a/src/engine/base.py +++ b/src/engine/base.py @@ -4,20 +4,21 @@ # ytdlbot - types.py import logging +import re import tempfile import uuid from abc import ABC, abstractmethod +from io import StringIO from pathlib import Path from types import SimpleNamespace -from tqdm import tqdm -from io import StringIO + import ffmpeg -import re import filetype from helper import debounce, sizeof_fmt from pyrogram import types +from tqdm import tqdm -from config import Types, TG_NORMAL_MAX_SIZE +from config import TG_NORMAL_MAX_SIZE, Types from database import Redis from database.model import ( get_download_settings, @@ -113,7 +114,7 @@ def more(title, initial): f.close() return text - def download_hook(self, d: dict, bot_msg): + def download_hook(self, d: dict): if d["status"] == "downloading": downloaded = d.get("downloaded_bytes", 0) total = d.get("total_bytes") or d.get("total_bytes_estimate", 0) @@ -127,15 +128,15 @@ def download_hook(self, d: dict, bot_msg): eta = self.__remove_bash_color(d.get("_eta_str", d.get("eta"))) text = self.__tqdm_progress("Downloading...", total, downloaded, speed, eta) # debounce in here - self.edit_text(bot_msg, text) + self.__edit_text(self._bot_msg, text) - def upload_hook(self, current, total, bot_msg): + def upload_hook(self, current, total): text = self.__tqdm_progress("Uploading...", total, current) - self.edit_text(bot_msg, text) + self.__edit_text(self._bot_msg, text) @debounce(5) - def edit_text(self, bot_msg: types.Message, text: str): - bot_msg.edit_text(text) + def __edit_text(self, text: str): + self._bot_msg.edit_text(text) def get_cache_fileid(self): unique = self._url + get_download_settings(self._url) @@ -170,7 +171,7 @@ def send_something(self, *, chat_id, files, _type, caption=None, thumb=None, **k files[0], caption=caption, thumb=thumb, - progress=upload_hook, + progress=self.upload_hook, progress_args=(self._bot_msg,), **kwargs, ) diff --git a/src/engine/direct.py b/src/engine/direct.py index 24fd51c0..4d7757e8 100644 --- a/src/engine/direct.py +++ b/src/engine/direct.py @@ -36,7 +36,7 @@ def _aria2_download(self): video_paths = None # Download process using aria2c try: - self._bot_msg.edit_text(f"Download Starting...", disable_web_page_preview=True) + self._bot_msg.__edit_text(f"Download Starting...", disable_web_page_preview=True) # Command to engine the link using aria2c command = [ "aria2c", @@ -58,14 +58,14 @@ def _aria2_download(self): line = process.stdout.readline().decode("utf-8") if line.startswith("[#"): line = line.strip() - self._bot_msg.edit_text(f"Downloading... \n\n`{line}`", disable_web_page_preview=True) + self._bot_msg.__edit_text(f"Downloading... \n\n`{line}`", disable_web_page_preview=True) break iteration += 1 if iteration >= max_iterations: - self._bot_msg.edit_text("Something went wrong. Please try again.", disable_web_page_preview=True) + self._bot_msg.__edit_text("Something went wrong. Please try again.", disable_web_page_preview=True) except Exception as e: - self._bot_msg.edit_text(f"Download failed!❌\n\n`{e}`", disable_web_page_preview=True) + self._bot_msg.__edit_text(f"Download failed!❌\n\n`{e}`", disable_web_page_preview=True) return # Get filename and extension correctly after engine filepath = list(pathlib.Path(tempdir).glob("*")) @@ -73,7 +73,7 @@ def _aria2_download(self): path_obj = pathlib.Path(file_path_obj) filename = path_obj.name logging.info("Downloaded file %s", filename) - self._bot_msg.edit_text(f"Download Complete", disable_web_page_preview=True) + self._bot_msg.__edit_text(f"Download Complete", disable_web_page_preview=True) ext = filetype.guess_extension(file_path_obj) # Rename file if it doesn't have extension if ext is not None and not filename.endswith(ext): @@ -83,7 +83,7 @@ def _aria2_download(self): video_paths = list(pathlib.Path(tempdir).glob("*")) self._client.send_chat_action(chat_id, enums.ChatAction.UPLOAD_DOCUMENT) upload_processor(self._client, self._bot_msg, self._url, video_paths) - self._bot_msg.edit_text("Download success!✅") + self._bot_msg.__edit_text("Download success!✅") def _download(self, formats): if ENABLE_ARIA2: diff --git a/src/engine/generic.py b/src/engine/generic.py index b09c542e..6c4801d5 100644 --- a/src/engine/generic.py +++ b/src/engine/generic.py @@ -32,7 +32,7 @@ def _setup_formats(self): # Add any remaining buttons as the last row if temp_row: markup.append(temp_row) - self._bot_msg.edit_text("Choose the format", reply_markup=types.InlineKeyboardMarkup(markup)) + self._bot_msg.__edit_text("Choose the format", reply_markup=types.InlineKeyboardMarkup(markup)) return if download == "audio": # download audio only