add visual progress bar

requirements.txt

@@ -1,6 +1,6 @@
 pyrogram==1.3.6
 tgcrypto==1.2.2
-yt-dlp==2021.12.27
+yt-dlp==2022.1.21
 APScheduler==3.8.1
 beautifultable==1.0.1
 ffmpeg-python==0.2.0

ytdlbot/constant.py

@@ -9,8 +9,8 @@ __author__ = "Benny <benny.think@gmail.com>"
 
 import time
 
-from config import (AFD_LINK, COFFEE_LINK, ENABLE_VIP, EX, MULTIPLY, OWNER,
-                    REQUIRED_MEMBERSHIP, USD2CNY)
+from config import (AFD_LINK, COFFEE_LINK, ENABLE_CELERY, ENABLE_VIP, EX,
+                    MULTIPLY, REQUIRED_MEMBERSHIP, USD2CNY)
 from downloader import sizeof_fmt
 from limit import QUOTA, VIP
 from utils import get_func_queue, get_queue_stat
@@ -118,10 +118,10 @@ Sending format: **{1}**
 
     @staticmethod
     def get_receive_link_text():
+        if ENABLE_CELERY:
             reserved = get_func_queue("reserved")
-        if reserved == 0:
-            text = "Your task was added to active queue.\nProcessing...\n\n"
-        else:
             text = f"Too many tasks. Your tasks was added to the reserved queue {reserved}."
+        else:
+            text = "Your task was added to active queue.\nProcessing...\n\n"
 
         return text

ytdlbot/downloader.py

@@ -10,14 +10,15 @@ __author__ = "Benny <benny.think@gmail.com>"
 import logging
 import os
 import pathlib
-import random
 import re
 import subprocess
 import time
+from io import StringIO
 
 import fakeredis
 import filetype
 import yt_dlp as ytdl
+from tqdm import tqdm
 from yt_dlp import DownloadError
 
 from config import ENABLE_VIP, TG_MAX_SIZE
@@ -42,12 +43,39 @@ def edit_text(bot_msg, text):
     key = f"{bot_msg.chat.id}-{bot_msg.message_id}"
     # if the key exists, we shouldn't send edit message
     if not r.exists(key):
-        r.set(key, "ok", ex=random.randint(1, 5))
+        r.set(key, "ok", ex=3)
         bot_msg.edit_text(text)
 
 
+def tqdm_progress(desc, total, finished, speed="", eta=""):
+    def more(title, initial):
+        if initial:
+            return f"{title} {initial}"
+        else:
+            return ""
+
+    f = StringIO()
+    tqdm(total=total, initial=finished, file=f, ascii=False, unit_scale=True, ncols=30,
+         bar_format="{l_bar}{bar} |{n_fmt}/{total_fmt} "
+         )
+    raw_output = f.getvalue()
+    tqdm_output = raw_output.split("|")
+    progress = f"[{tqdm_output[1]}]"
+    detail = tqdm_output[2]
+    text = f"""
+{desc}
+
+{progress}
+{detail}
+{more("Speed:", speed)}
+{more("ETA:", eta)}
+    """
+    f.close()
+    return text
+
+
 def remove_bash_color(text):
-    return re.sub(r'\u001b|\[0;94m|\u001b\[0m|\[0;32m|\[0m', "", text)
+    return re.sub(r'\u001b|\[0;94m|\u001b\[0m|\[0;32m|\[0m|\[0;33m', "", text)
 
 
 def download_hook(d: dict, bot_msg):
@@ -60,7 +88,6 @@ def download_hook(d: dict, bot_msg):
     if d['status'] == 'downloading':
         downloaded = d.get("downloaded_bytes", 0)
         total = d.get("total_bytes") or d.get("total_bytes_estimate", 0)
-        # total = 0
         filesize = sizeof_fmt(total)
         max_size = 2 * 1024 * 1024 * 1024
         if total > max_size:
@@ -74,14 +101,15 @@ def download_hook(d: dict, bot_msg):
             result, err_msg = check_quota(total, bot_msg.chat.id)
             if result is False:
                 raise ValueError(err_msg)
-        text = f'[{filesize}]: Downloading {percent} - {downloaded}/{total} @ {speed}'
+        eta = remove_bash_color(d.get("_eta_str", d.get("eta")))
+        text = tqdm_progress("Downloading...", total, downloaded, speed, eta)
         edit_text(bot_msg, text)
         r.set(key, "ok", ex=5)
 
 
 def upload_hook(current, total, bot_msg):
-    filesize = sizeof_fmt(total)
-    text = f'[{filesize}]: Uploading {round(current / total * 100, 2)}% - {current}/{total}'
+    # filesize = sizeof_fmt(total)
+    text = tqdm_progress("Uploading...", total, current)
     edit_text(bot_msg, text)
 
 

ytdlbot/tasks.py

@@ -17,16 +17,17 @@ import time
 from urllib.parse import quote_plus
 
 import requests
+from apscheduler.schedulers.background import BackgroundScheduler
 from celery import Celery
 from pyrogram import idle
 from pyrogram.types import InlineKeyboardButton, InlineKeyboardMarkup
-from apscheduler.schedulers.background import BackgroundScheduler
 
 from client_init import create_app
 from config import BROKER, ENABLE_CELERY, ENABLE_VIP, OWNER, WORKERS
 from constant import BotText
 from db import Redis
-from downloader import convert_flac, sizeof_fmt, upload_hook, ytdl_download
+from downloader import (convert_flac, edit_text, sizeof_fmt, tqdm_progress,
+                        upload_hook, ytdl_download)
 from limit import VIP
 from utils import (apply_log_formatter, auto_restart, customize_logger,
                    get_metadata, get_user_settings)
@@ -104,7 +105,7 @@ def direct_normal_download(bot_msg, client, url):
     headers = {
         "user-agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/80.0.3987.149 Safari/537.36"}
     vip = VIP()
-
+    length = 0
     if ENABLE_VIP:
         remain, _, _ = vip.check_remaining_quota(chat_id)
         try:
@@ -116,9 +117,10 @@ def direct_normal_download(bot_msg, client, url):
             bot_msg.reply_text(f"Sorry, you have reached your quota.\n")
             return
 
-    req = ""
+    req = None
     try:
-        req = requests.get(url, headers=headers)
+        req = requests.get(url, headers=headers, stream=True)
+        length = int(req.headers.get("content-length"))
         filename = re.findall("filename=(.+)", req.headers.get("content-disposition"))[0]
     except TypeError:
         filename = getattr(req, "url", "").rsplit("/")[-1]
@@ -131,9 +133,14 @@ def direct_normal_download(bot_msg, client, url):
 
     with tempfile.TemporaryDirectory() as f:
         filepath = f"{f}/{filename}"
-        print(filepath)
-        with open(filepath, "wb") as fp:
-            fp.write(req.content)
+        # consume the req.content
+        downloaded = 0
+        for chunk in req.iter_content(1024 * 1024):
+            text = tqdm_progress("Downloading...", length, downloaded)
+            edit_text(bot_msg, text)
+            with open(filepath, "ab") as fp:
+                fp.write(chunk)
+            downloaded += len(chunk)
         logging.info("Downloaded file %s", filename)
         st_size = os.stat(filepath).st_size
         if ENABLE_VIP:
@@ -167,10 +174,14 @@ def normal_audio(bot_msg):
 
 def get_worker_status(username):
     worker_name = os.getenv("WORKER_NAME")
+    try:
         me = celery_client.get_me()
+        mention = me.mention()
+    except Exception:
+        mention = "YouTube Downloader"
     if worker_name and username == OWNER:
-        return f"Downloaded by {me.mention()}-{worker_name}"
-    return f"Downloaded by {me.mention()}"
+        return f"Downloaded by {mention}-{worker_name}"
+    return f"Downloaded by {mention}"
 
 
 def ytdl_normal_download(bot_msg, client, url):