Add queue system

This commit is contained in:
Nexus 2023-05-02 19:38:39 +01:00
parent 1b0b1630e0
commit 433febae89
Signed by: nex
GPG key ID: 0FA334385D0B689F
2 changed files with 187 additions and 42 deletions

View file

@ -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.")

View file

@ -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"
]