From 779a44d9e7efa848d5509c825949cadf128d578f Mon Sep 17 00:00:00 2001 From: BennyThink Date: Tue, 25 Jan 2022 18:44:50 +0800 Subject: [PATCH] add visual progress bar --- requirements.txt | 2 +- ytdlbot/constant.py | 12 ++++++------ ytdlbot/downloader.py | 42 +++++++++++++++++++++++++++++++++++------- ytdlbot/tasks.py | 33 ++++++++++++++++++++++----------- 4 files changed, 64 insertions(+), 25 deletions(-) diff --git a/requirements.txt b/requirements.txt index a3734e9..a4ebbd6 100644 --- a/requirements.txt +++ b/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 diff --git a/ytdlbot/constant.py b/ytdlbot/constant.py index 3e74cc8..3e1fe9b 100644 --- a/ytdlbot/constant.py +++ b/ytdlbot/constant.py @@ -9,8 +9,8 @@ __author__ = "Benny " 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(): - reserved = get_func_queue("reserved") - if reserved == 0: - text = "Your task was added to active queue.\nProcessing...\n\n" - else: + if ENABLE_CELERY: + reserved = get_func_queue("reserved") 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 diff --git a/ytdlbot/downloader.py b/ytdlbot/downloader.py index 17eabb8..5103cf0 100644 --- a/ytdlbot/downloader.py +++ b/ytdlbot/downloader.py @@ -10,14 +10,15 @@ __author__ = "Benny " 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) diff --git a/ytdlbot/tasks.py b/ytdlbot/tasks.py index 986e773..5397f7e 100644 --- a/ytdlbot/tasks.py +++ b/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") - me = celery_client.get_me() + 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):