diff --git a/src/cogs/ytdl.py b/src/cogs/ytdl.py index 6cdd05e..1fc5144 100644 --- a/src/cogs/ytdl.py +++ b/src/cogs/ytdl.py @@ -2,6 +2,9 @@ import asyncio import functools import hashlib import logging +import math +import time + import httpx import subprocess import tempfile @@ -257,7 +260,32 @@ class YTDLCog(commands.Cog): ): """Runs yt-dlp and outputs into discord.""" await ctx.defer() + last_edit = time.time() options = self.default_options.copy() + + async def _download_hook(_data: dict[str, typing.Any]): + n = time.time() + _total = _data.get("total_bytes", _data.get("total_bytes_estimate")) + if _total: + _percent = round(_data.get("downloaded_bytes", 0) / _total * 100, 2) + else: + _total = max(1, _data.get("fragment_count", 4096)) + _percent = round(max(_data.get("fragment_index", 1), 1) / _total * 100, 2) + _speed_bytes_per_second = _data.get("speed", 1) + _speed_megabits_per_second = round((_speed_bytes_per_second * 8) / 1024 / 1024) + _eta = discord.utils.utcnow() + _data.get("eta") or discord.utils.utcnow() + blocks = "#" * math.floor(_percent / 10) + bar = f"{blocks}{'.' * (10 - len(blocks))}" + line = (f"{_percent}% [{bar}] | {_speed_megabits_per_second}Mbps | " + f"ETA {discord.utils.format_dt(_eta, 'R')}") + nonlocal last_edit + if (n - last_edit) >= 5: + embed.clear_fields() + embed.add_field(name="Progress", value=line) + await ctx.edit(embed=embed) + last_edit = time.time() + options["progress_hooks"] = [_download_hook] + description = "" with tempfile.TemporaryDirectory(prefix="jimmy-ytdl-") as temp_dir: @@ -355,15 +383,20 @@ class YTDLCog(commands.Cog): description += "\n".join(lines) domain = urlparse(webpage_url).netloc + embed = discord.Embed( + title=title, + description=description, + url=webpage_url, + colour=self.colours.get(domain, discord.Colour.og_blurple()), + ) + embed.add_field( + name="Progress", + value="0% [..........] " + ) + embed.set_footer(text="Downloading (step 2/10)") + embed.set_thumbnail(url=thumbnail_url) await ctx.edit( - embed=discord.Embed( - title=title, - description=description, - url=webpage_url, - colour=self.colours.get(domain, discord.Colour.og_blurple()), - ) - .set_footer(text="Downloading (step 2/10)") - .set_thumbnail(url=thumbnail_url) + embed=embed ) previous = await self.get_saved(webpage_url, chosen_format_id, snip or "*") if previous: @@ -379,6 +412,9 @@ class YTDLCog(commands.Cog): ).set_image(url=previous), ) return + + last_edit = time.time() + try: await asyncio.to_thread(functools.partial(downloader.download, [url])) except yt_dlp.DownloadError as e: