From 433febae8962d90e03c63cbc0a1f2bb375e15226 Mon Sep 17 00:00:00 2001 From: nex Date: Tue, 2 May 2023 19:38:39 +0100 Subject: [PATCH] Add queue system --- cogs/voice.py | 223 +++++++++++++++++++++++++++++++++++++++--------- utils/client.py | 6 +- 2 files changed, 187 insertions(+), 42 deletions(-) diff --git a/cogs/voice.py b/cogs/voice.py index dc17a5f..c8cbd1f 100644 --- a/cogs/voice.py +++ b/cogs/voice.py @@ -1,10 +1,11 @@ +import functools import io +import json import shutil import asyncio import discord import yt_dlp import tempfile -from datetime import datetime from pathlib import Path from discord.ext import commands @@ -18,8 +19,8 @@ class TransparentQueue(asyncio.Queue): await super().put(item) self._internal_queue.append(item) - async def job_done(self): - await super().job_done() + def task_done(self): + super().task_done() self._internal_queue.pop(0) @@ -58,9 +59,10 @@ class YTDLSource(discord.PCMVolumeTransformer): class VoiceCog(commands.Cog): def __init__(self, bot): self.bot = bot + self.cache = Path(tempfile.mkdtemp("jimmy-voice-cache-")).resolve() + self.cache.mkdir(exist_ok=True, parents=True) self.ytdl_options = { "format": "bestaudio/best", - "outtmpl": "%(title)s.%(ext)s", "restrictfilenames": True, "noplaylist": True, "quiet": True, @@ -68,40 +70,59 @@ class VoiceCog(commands.Cog): "ignoreerrors": True, "logtostderr": False, "default_search": "auto", + "outtmpl": "%(title)s.%(ext)s".format(self.cache), + "keepvideo": False, + "cachedir": str(self.cache), + "paths": { + "home": str(self.cache), + "temp": str(self.cache), + } } - self.cache = Path(tempfile.mkdtemp("jimmy-voice-cache-")).resolve() self.yt_dl = yt_dlp.YoutubeDL(self.ytdl_options) - # self.queue = TransparentQueue(100) - # self._queue_task = self.bot.loop.create_task(self.queue_task()) + self.queue = TransparentQueue(100) + self._queue_task = self.bot.loop.create_task(self.queue_task()) + self.song_done = asyncio.Event() async def queue_task(self): if not self.bot.is_ready(): await self.bot.wait_until_ready() while True: - ctx, player, author, inserted_at = await self.queue.get() + ctx, player, inserted_at = await self.queue.get() + if not ctx.guild.voice_client: + # no longer playing. clear queue + del self.queue + self.queue = TransparentQueue(100) + continue ctx.guild.voice_client.play(player, after=self.after_player(ctx)) + self.song_done.clear() embed = discord.Embed( - description=f"Now playing: [{player.title}]({player.url}), as requested by {author.mention}.", + description=f"Now playing: [{player.title}]({player.url}), as requested by {ctx.author.mention}, " + f"{discord.utils.format_dt(inserted_at, 'R')}.", color=discord.Color.green(), ) try: - await ctx.respond( - embed=embed + await ctx.guild.voice_client.channel.send( + ctx.author.mention, + embed=embed, + allowed_mentions=discord.AllowedMentions(users=True), ) except discord.HTTPException: try: - await ctx.send(embed=embed) + await ctx.respond(embed=embed) except discord.HTTPException: pass - - await self.queue.job_done() + self.queue.task_done() + await self.song_done.wait() def cog_unload(self): + for task in (self._queue_task,): + task.cancel() shutil.rmtree(self.cache) def after_player(self, ctx: discord.ApplicationContext): def after(e): + self.song_done.set() if e: self.bot.loop.create_task( ctx.respond( @@ -111,29 +132,70 @@ class VoiceCog(commands.Cog): return after async def unblock(self, func, *args, **kwargs): - return await self.bot.loop.run_in_executor(None, func, *args, **kwargs) + call = functools.partial(func, *args, **kwargs) + return await self.bot.loop.run_in_executor(None, call) @commands.slash_command(name="play") - async def play(self, ctx: discord.ApplicationContext, url: str, volume: float = 50): + async def stream(self, ctx: discord.ApplicationContext, url: str, volume: float = 100): """Streams a URL using yt-dl""" if not ctx.user.voice: await ctx.respond("You are not connected to a voice channel.") return - if ctx.voice_client.is_playing(): - await ctx.respond("Already playing audio. Use %s first." % self.bot.get_application_command("stop").mention) - return - player = await YTDLSource.from_url(self.yt_dl, url, loop=self.bot.loop, stream=True) if not player: await ctx.respond("Could not extract any audio from the given URL.") - ctx.guild.voice_client.play(player, after=self.after_player(ctx)) - ctx.guild.voice_client.source.volume = min(100.0, max(1.0, volume / 100)) + + player.volume = volume / 100 + await self.queue.put((ctx, player, discord.utils.utcnow())) embed = discord.Embed( - description=f"Playing [{player.title}]({player.url})", + description=f"Added [{player.title}]({player.url}) to the queue.", ) await ctx.respond(embed=embed) + @commands.command(hidden=True) + async def play(self, ctx: commands.Context, url: str, volume: float = 100): + """Plays a song by downloading it first.""" + if not ctx.author.voice: + await ctx.reply("You are not connected to a voice channel.") + return + + player = await YTDLSource.from_url(self.yt_dl, url, loop=self.bot.loop) + if not player: + await ctx.reply("Could not extract any audio from the given URL.") + + player.volume = volume / 100 + await self.queue.put((ctx, player, discord.utils.utcnow())) + embed = discord.Embed( + description=f"Added [{player.title}]({player.url}) to the queue.", + ) + await ctx.reply(embed=embed) + + @commands.slash_command(name="queue") + async def view_queue(self, ctx: discord.ApplicationContext): + """Views the current queue""" + if not ctx.guild.voice_client: + await ctx.respond("Not connected to a voice channel.") + return + + if not self.queue._internal_queue and not ctx.guild.voice_client.is_playing(): + await ctx.respond("The queue is empty.") + return + + embed = discord.Embed( + title="Queue", + description="\n".join( + f"{i+1}. [{x[1].title}]({x[1].url})" for i, x in enumerate(self.queue._internal_queue) + ), + ) + if ctx.guild.voice_client.is_playing(): + now_playing = ctx.guild.voice_client.source + embed.add_field( + name="Now Playing", + value=f"[{now_playing.title}]({now_playing.url})", + ) + await ctx.respond(embed=embed) + @commands.slash_command(name="volume") async def volume(self, ctx: discord.ApplicationContext, volume: float): """Changes the player's volume""" @@ -153,41 +215,122 @@ class VoiceCog(commands.Cog): if ctx.voice_client: if ctx.voice_client.is_playing(): - # members = ctx.voice_client.channel.members - # bots = [m for m in members if m.bot] - # if len(bots) == len(members): - # pass - # else: - # humans = len(members) - len(bots) - # if humans > 1: - # ctx.voice_client.stop() await ctx.voice_client.disconnect(force=True) await ctx.respond("Disconnected from voice channel.") else: await ctx.respond("Not connected to a voice channel.") + @commands.slash_command(name="skip") + async def skip(self, ctx: discord.ApplicationContext): + """Skips the current song""" + class VoteSkipDialog(discord.ui.View): + def __init__(self): + super().__init__() + self.voted = [] + + @discord.ui.button(label="Skip", style=discord.ButtonStyle.green) + async def _skip(self, button: discord.ui.Button, interaction: discord.Interaction): + await interaction.response.defer(ephemeral=True) + humans = len(ctx.guild.voice_client.channel.members) - len( + [x for x in ctx.guild.voice_client.channel.members if x.bot] + ) + target = humans / 2 if humans > 2 else humans + if interaction.user.id not in self.voted: + self.voted.append(interaction.user.id) + if humans <= 1 or target <= 1: + ctx.voice_client.stop() + self.stop() + self.disable_all_items() + return await interaction.edit_original_response( + view=self, + content="Skipped song." + ) + + if len(self.voted) >= target: + ctx.voice_client.stop() + self.stop() + self.disable_all_items() + return await interaction.edit_original_response( + view=self, + content="Skipped song." + ) + else: + await ctx.respond( + f"Voted to skip. %d/%d" % (len(self.voted), target), + ) + else: + await ctx.respond("You have already voted to skip.") + + self.disable_all_items() + await interaction.edit_original_response( + view=self, + content="Vote skip (%d/%d)." % (len(self.voted), target) + ) + + if not ctx.guild.voice_client: + await ctx.respond("Not connected to a voice channel.") + return + + if ctx.voice_client.is_playing(): + _humans = len(ctx.guild.voice_client.channel.members) - len( + [x for x in ctx.guild.voice_client.channel.members if x.bot] + ) + if _humans > 1: + _target = _humans / 2 if _humans > 2 else _humans + diag = VoteSkipDialog() + diag.voted.append(ctx.user.id) + await ctx.respond("Vote skip (1/%d)." % _target, view=VoteSkipDialog()) + ctx.voice_client.stop() + self.song_done.set() + await ctx.respond("Skipped song.") + else: + await ctx.respond("Not playing any music.") + @commands.command(name="dump-metadata") - async def dump_metadata(self, ctx: commands.Context, *, url: str): + async def dump_metadata(self, ctx: commands.Context, url: str, traverse: str = None): """Dumps JSON YT-DLP metadata to a file""" + content = None + _params = self.yt_dl.params.copy() + _params.pop("noplaylist") + _params.pop("logtostderr") + _ytdl = yt_dlp.YoutubeDL(_params) async with ctx.channel.typing(): file = io.StringIO() - data = await self.unblock(self.yt_dl.extract_info, url, download=False) - data = await self.unblock(self.yt_dl.sanitize_info, data) + data = await self.unblock(_ytdl.extract_info, url) or {} + data = await self.unblock(_ytdl.sanitize_info, data, remove_private_keys=True) + if traverse: + last_key = [] + for key in traverse.split("."): + if key in data: + data = data[key] + last_key.append(key) + else: + content = "Key %r not found in metadata (got as far as %r)." % (key, '.'.join(last_key)) + break json.dump(data, file, indent=4) file.seek(0) - return await ctx.respond(file=discord.File(file, filename="metadata.json")) + # noinspection PyTypeChecker + return await ctx.reply(content, file=discord.File(file, filename="metadata.json")) - async def cog_before_invoke(self, ctx: discord.ApplicationContext): - await ctx.defer() + @volume.before_invoke + @play.before_invoke + @stop.before_invoke + @stream.before_invoke + async def before_invoke(self, ctx: discord.ApplicationContext | commands.Context): + if isinstance(ctx, discord.ApplicationContext): + await ctx.defer() + sender = ctx.respond + else: + sender = ctx.reply if not self.cache.exists(): self.cache.mkdir() if not ctx.guild.voice_client: - if ctx.user.voice: - await ctx.user.voice.channel.connect() + if ctx.author.voice: + await ctx.author.voice.channel.connect() else: - await ctx.respond("You are not connected to a voice channel.") + await sender("You are not connected to a voice channel.") raise commands.CommandError("User not connected to a voice channel.") diff --git a/utils/client.py b/utils/client.py index 14a775d..3646e58 100644 --- a/utils/client.py +++ b/utils/client.py @@ -38,7 +38,9 @@ class Bot(commands.Bot): for ext in extensions: try: self.load_extension(ext) - except (discord.ExtensionNotFound, discord.ExtensionFailed, OSError) as e: + except discord.ExtensionNotFound: + console.log(f"[red]Failed to load extension {ext}: Extension not found.") + except (discord.ExtensionFailed, OSError) as e: console.log(f"[red]Failed to load extension {ext}: {e}") if getattr(config, "dev", False): console.print_exception() @@ -98,7 +100,7 @@ except ImportError: "cogs.other", "cogs.starboard", "cogs.uptime", -# "cogs.voice" + "cogs.voice" ]