mirror of
https://github.com/tgbot-collection/ytdlbot.git
synced 2024-09-20 23:15:57 +08:00
add visual progress bar
This commit is contained in:
parent
e915a33219
commit
779a44d9e7
|
@ -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
|
||||
|
|
|
@ -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():
|
||||
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
|
||||
|
|
|
@ -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)
|
||||
|
||||
|
||||
|
|
|
@ -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):
|
||||
|
|
Loading…
Reference in a new issue