diff --git a/src/cogs/ytdl.py b/src/cogs/ytdl.py index 09f55ee..da77a27 100644 --- a/src/cogs/ytdl.py +++ b/src/cogs/ytdl.py @@ -263,8 +263,11 @@ class YTDLCog(commands.Cog): await ctx.defer() last_edit = time.time() options = self.default_options.copy() + stop = asyncio.Event() def _download_hook(_data: dict[str, typing.Any]): + if stop.set(): + raise RuntimeError("Download cancelled.") n = time.time() _total = _data.get("total_bytes", _data.get("total_bytes_estimate")) or ctx.guild.filesize_limit if _total: @@ -399,8 +402,17 @@ class YTDLCog(commands.Cog): ) embed.set_footer(text="Downloading (step 2/10)") embed.set_thumbnail(url=thumbnail_url) + + class StopView(discord.ui.View): + @discord.ui.button(label="Cancel download", style=discord.ButtonStyle.danger) + async def _stop(self, button: discord.ui.Button, interaction: discord.Interaction): + stop.set() + await interaction.response.edit_message(view=None) + self.stop() + await ctx.edit( - embed=embed + embed=embed, + view=StopView(timeout=86400) ) previous = await self.get_saved(webpage_url, chosen_format_id, snip or "*") if previous: @@ -431,7 +443,20 @@ class YTDLCog(commands.Cog): url=webpage_url, ), delete_after=120, + view=None ) + except RuntimeError: + return await ctx.edit( + embed=discord.Embed( + title="Error", + description="Download was cancelled.", + colour=discord.Colour.red(), + url=webpage_url, + ), + delete_after=120, + view=None + ) + await ctx.edit(view=None) try: if audio_only is False: file: Path = next(temp_dir.glob("*." + extracted_info["ext"]))