Partial migration to dpy, music cog

This commit is contained in:
2022-11-12 19:38:22 +01:00
parent a4986a04fc
commit 9f62a48bbd
15 changed files with 446 additions and 611 deletions
+196
View File
@@ -0,0 +1,196 @@
import typing
from cogs.music.commands import MusicCommands
import asyncio
from typing import Optional
import lavalink
from discord import Embed
from discord.channel import TextChannel
from discord.ext import commands
from discord.ext.commands.errors import CommandError
from lavalink.models import AudioTrack
from lavalink.models import DefaultPlayer
from cogs.music import helper
from bot import TuneBot
from cogs.music.voice_client import LavalinkVoiceClient
from utils.embed import create_embed
from utils.log import logger
from context import CustomContext
from utils.classes import BaseCog
from utils.EmbedGenerator import EmbedGenerator
from utils.exceptions import EmbeddedCommandException
if typing.TYPE_CHECKING:
from bot import TuneBot
class MusicCog(BaseCog):
def __init__(self, bot: TuneBot):
super().__init__(bot)
if not hasattr(bot, "lavalink"):
self.bot.lavalink = self.bot.create_lavalink(bot.user.id)
self.bot.lavalink.add_event_hook(self.track_hook)
@commands.Cog.listener()
async def on_ready(self):
while not self.is_lavalink_ready():
await asyncio.sleep(1)
redis_result = await self.bot.global_autojoin.fetch_channels()
for guild_id, (voicechannel_id, textchannel_id) in redis_result.items():
player = self.bot.lavalink.player_manager.create(guild_id)
player.store("channel", textchannel_id)
voice_channel = await self.bot.fetch_channel(voicechannel_id)
await voice_channel.connect(cls=LavalinkVoiceClient)
if not player.is_playing:
await self.fill_player_queue(
player, self.bot.config["queue_buffer_size"]
)
await player.play()
async def fill_player_queue(self, player: DefaultPlayer, amount: Optional[int] = 1):
await helper.fill_player_queue(self.bot, player, amount)
def cog_unload(self):
"""Cog unload handler. This removes any event hooks that were registered."""
self.bot.lavalink._event_hooks.clear()
async def cog_before_invoke(self, ctx: CustomContext):
"""Command before-invoke handler."""
guild_check = ctx.guild is not None
# This is essentially the same as `@commands.guild_only()`
# except it saves us repeating ourselves (and also a few lines).
if guild_check:
await self.ensure_voice(ctx)
# Ensure that the bot and command author share a mutual voicechannel.
return guild_check
async def cog_command_error(self, ctx: CustomContext, error: CommandError):
if isinstance(error, commands.CommandInvokeError):
await ctx.send(error.original)
elif isinstance(error, EmbeddedCommandException):
await error.send(ctx)
async def ensure_voice(self, ctx: CustomContext):
"""This check ensures that the bot and command author are in the same voicechannel."""
player = self.bot.lavalink.player_manager.create(ctx.guild.id)
# Create returns a player if one exists, otherwise creates.
# This line is important because it ensures that a player always exists for a guild.
# Most people might consider this a waste of resources for guilds that aren't playing, but this is
# the easiest and simplest way of ensuring players are created.
# These are commands that require the bot to join a voicechannel (i.e. initiating playback).
# Commands such as volume/skip etc don't require the bot to be in a voicechannel so don't need listing here.
should_connect = ctx.command.name in ("connect",)
if should_connect:
try:
await helper.ensure_voice(
permissions=ctx.me.guild_permissions,
player=player,
author=ctx.author,
voice_client=ctx.guild.voice_client,
channel_id=ctx.channel.id,
)
except Exception as exc:
raise commands.CommandInvokeError(exc)
async def track_hook(self, event: lavalink.Event):
if isinstance(event, lavalink.events.QueueEndEvent):
# When this track_hook receives a "QueueEndEvent" from lavalink.py
# it indicates that there are no tracks left in the player's queue.
# To save on resources, we can tell the bot to disconnect from the voicechannel.
guild_id = event.player.guild_id
guild = self.bot.get_guild(guild_id)
await guild.voice_client.disconnect(force=True)
elif isinstance(event, lavalink.events.TrackStartEvent):
channel_id = int(event.player.fetch("channel"))
channel: TextChannel = self.bot.get_channel(channel_id)
embed = helper.create_track_embed(event.player.current)
await channel.send(embed=embed)
elif isinstance(event, lavalink.events.TrackEndEvent):
await self.fill_player_queue(event.player, 1)
@commands.command(name="connect", aliases=["p", "play", "join"])
async def play(self, ctx: CustomContext):
"""Start the radio"""
logger.info(f"Joining channel {ctx.channel}")
# Get the player for this guild from cache.
if ctx.player.is_connected:
await ctx.send("Already connected")
return
await self.fill_player_queue(
ctx.player, self.bot.config["queue_buffer_size"] + 1
)
if not ctx.player.is_playing:
await ctx.player.play()
await EmbedGenerator.Title(ctx, "*⃣ | Connected.")
return
@commands.command(name="skip", aliases=["next"])
async def skip(self, ctx: CustomContext):
"""Skip the current song"""
await ctx.player.skip()
await ctx.send("Skipped current song")
@commands.command(name="queue")
async def queue(self, ctx: CustomContext):
"""Display the current radio queue"""
embed_color = self.bot.colors["embed"]
embed = Embed(title="Coming Up...", colour=embed_color)
if len(ctx.player.queue) > 0:
embed.description = "\n".join(
f"{index}. [{track.title}]({track.uri})"
for index, track in enumerate(ctx.player.queue, 1)
)
else:
embed.description = "We are still determining a playlist"
await ctx.send(embed=embed)
@commands.command(name="disconnect", aliases=["dc", "stop"])
async def disconnect(self, ctx: CustomContext):
"""Disconnects the radio from the channel"""
if not ctx.author.voice or (
ctx.player.is_connected
and ctx.author.voice.channel.id != int(ctx.player.channel_id)
):
# Abuse prevention. Users not in voice channels, or not in the same voice channel as the bot
# may not disconnect the bot.
await EmbedGenerator.Title(ctx, "You're not in my voicechannel!")
return
# Clear the queue to ensure old tracks don't start playing
# when someone else queues something.
ctx.player.queue.clear()
# Stop the current track so Lavalink consumes less resources.
await ctx.player.stop()
# Disconnect from the voice channel.
await ctx.voice_client.disconnect(force=True)
await EmbedGenerator.Title(ctx, "*⃣ | Disconnected.")
@commands.command(name="now")
async def now_playing(self, ctx: CustomContext):
"""Displays information about the currently played track"""
if not ctx.player.current:
embed = create_embed(ctx.author)
embed.title = "You're not listening to anything right now."
await ctx.send(embed=embed)
return
embed = helper.create_track_embed(ctx.player.current)
await ctx.send(embed=embed)
async def setup(bot: "TuneBot"):
await bot.add_cog(MusicCog(bot))
bot.tree.add_command(MusicCommands(bot), override=True)
+110
View File
@@ -0,0 +1,110 @@
import typing
from discord import Interaction, app_commands
from bot import config
from utils.embed import create_embed
from cogs.music import helper
if typing.TYPE_CHECKING:
from bot import TuneBot
QUEUE_SIZE = config["queue_buffer_size"] + 1
class MusicCommands(app_commands.Group):
def __init__(self, bot: "TuneBot"):
super().__init__(name="radio", description="Never-ending stream of good vibes")
self.bot: "TuneBot" = bot
@app_commands.command(name="connect", description="Start the radio")
async def connect(self, ctx: Interaction):
embed = create_embed(ctx.user)
player = helper.get_player(ctx.client, ctx.guild_id)
try:
await helper.ensure_voice(
permissions=ctx.app_permissions,
player=player,
author=ctx.user,
voice_client=ctx.guild.voice_client,
channel_id=ctx.channel.id,
)
except Exception as exc:
embed.title = str(exc)
await ctx.response.send_message(embed=embed)
return
if player.is_connected:
embed.title = "Already connected."
await ctx.response.send_message(embed=embed)
return
await helper.fill_player_queue(ctx.client, player, QUEUE_SIZE)
if not player.is_playing:
await player.play()
embed.title = "*⃣ | Connected."
await ctx.response.send_message(embed=embed)
@app_commands.command(
name="disconnect", description="I've had enough music for a while"
)
async def disconnect(self, ctx: Interaction):
embed = create_embed(ctx.user)
player = helper.get_player(ctx.client, ctx.guild_id)
if not ctx.user.voice or (
player.is_connected and ctx.user.voice.channel.id != int(player.channel_id)
):
embed.title = "You're not in my voicechannel!"
await ctx.response.send_message(embed=embed)
return
# Clear the queue to ensure old tracks don't start playing
# when someone else queues something.
player.queue.clear()
# Stop the current track so Lavalink consumes less resources.
await player.stop()
# Disconnect from the voice channel.
await ctx.guild.voice_client.disconnect(force=True)
embed.title = "*⃣ | Disconnected."
await ctx.response.send_message(embed=embed)
@app_commands.command(
name="current", description="Gives more details about what's currently playing."
)
async def current(self, ctx: Interaction):
player = helper.get_player(ctx.client, ctx.guild_id)
if not player.current:
embed = create_embed(ctx.user)
embed.title = "You're not listening to anything right now."
await ctx.response.send_message(embed=embed)
return
embed = helper.create_track_embed(player.current)
await ctx.response.send_message(embed=embed)
@app_commands.command(name="queue", description="See what's ahead")
async def queue(self, ctx: Interaction):
player = helper.get_player(ctx.client, ctx.guild_id)
embed = create_embed(ctx.user)
embed.title = "Coming up..."
if len(player.queue) > 0:
embed.description = "\n".join(
f"{index}. [{track.title}]({track.uri})"
for index, track in enumerate(player.queue, 1)
)
else:
embed.description = "We are still determining a playlist"
await ctx.response.send_message(embed=embed)
@app_commands.command(
name="skip", description="You might like the next track better"
)
async def skip(self, ctx: Interaction):
player = helper.get_player(ctx.client, ctx.guild_id)
embed = create_embed(ctx.user)
embed.title = "Skipped current song"
await player.skip()
await ctx.response.send_message(embed=embed)
+79
View File
@@ -0,0 +1,79 @@
import datetime
import typing
import discord
import lavalink
from cogs.music.voice_client import LavalinkVoiceClient
from utils.log import logger
from bot import colors
if typing.TYPE_CHECKING:
from lavalink import DefaultPlayer
from bot import TuneBot
def get_player(bot: "TuneBot", guild_id: int) -> "DefaultPlayer":
if player := bot.lavalink.player_manager.get(guild_id=guild_id):
return player
return bot.lavalink.player_manager.create(guild_id)
async def fill_player_queue(
bot: "TuneBot", player: "DefaultPlayer", amount: typing.Optional[int] = 1
):
queries = await bot.global_playlist.pick_random(amount)
logger.info(f"Got queries for {player.guild_id} {queries=} {amount=}")
for query in queries:
result = await player.node.get_tracks(query)
if not result or not result["tracks"]:
continue
track = lavalink.models.AudioTrack(
result["tracks"][0], bot.user.id, recommended=False
)
player.add(requester=bot.user.id, track=track)
def create_track_embed(track: lavalink.models.AudioTrack) -> discord.Embed:
embed = discord.Embed(
title="Now playing...",
colour=colors["embed"],
)
embed.description = f"[{track.title}]({track.uri})"
embed.set_thumbnail(url=f"https://i3.ytimg.com/vi/{track.identifier}/mqdefault.jpg")
try:
duration = str(datetime.timedelta(milliseconds=int(track.duration)))
except OverflowError:
duration = "🔴 LIVE"
embed.add_field(name="Duration", value=duration)
embed.add_field(name="Author", value=track.author)
return embed
async def ensure_voice(
permissions: discord.Permissions,
player: "DefaultPlayer",
author: discord.Member,
voice_client: typing.Any,
channel_id: int,
):
"""This ensures that the bot and command author are in the same voicechannel."""
if not author.voice or not author.voice.channel:
raise Exception("Join a voicechannel first.")
if not voice_client:
if not permissions.connect or not permissions.speak: # Check user limit too?
raise Exception("I need to be able to join your channel and speak")
player.store("channel", channel_id)
await author.voice.channel.connect(cls=LavalinkVoiceClient)
else:
if voice_client.channel.id != author.voice.channel.id:
raise Exception("You need to be in my voicechannel")
__all__ = ("fill_player_queue", "create_track_embed", "get_player", "ensure_voice")
+63
View File
@@ -0,0 +1,63 @@
import typing
import discord
if typing.TYPE_CHECKING:
from bot import TuneBot
class LavalinkVoiceClient(discord.VoiceClient):
def __init__(self, client: "TuneBot", channel: discord.abc.Connectable):
self.client = client
self.channel = channel
self.lavalink = self.client.lavalink
async def on_voice_server_update(self, data):
# the data needs to be transformed before being handed down to
# voice_update_handler
lavalink_data = {"t": "VOICE_SERVER_UPDATE", "d": data}
await self.lavalink.voice_update_handler(lavalink_data)
async def on_voice_state_update(self, data):
# the data needs to be transformed before being handed down to
# voice_update_handler
lavalink_data = {"t": "VOICE_STATE_UPDATE", "d": data}
await self.lavalink.voice_update_handler(lavalink_data)
async def connect(
self,
*,
timeout: float,
reconnect: bool,
self_deaf: bool = False,
self_mute: bool = False,
) -> None:
"""
Connect the bot to the voice channel and create a player_manager
if it doesn't exist yet.
"""
# ensure there is a player_manager when creating a new voice_client
self.lavalink.player_manager.create(guild_id=self.channel.guild.id)
await self.channel.guild.change_voice_state(
channel=self.channel, self_mute=self_mute, self_deaf=self_deaf
)
async def disconnect(self, *, force: bool) -> None:
"""
Handles the disconnect.
Cleans up running player and leaves the voice client.
"""
player = self.lavalink.player_manager.get(self.channel.guild.id)
# no need to disconnect if we are not connected
if not force and not player.is_connected:
return
# None means disconnect
await self.channel.guild.change_voice_state(channel=None)
# update the channel_id of the player to None
# this must be done because the on_voice_state_update that
# would set channel_id to None doesn't get dispatched after the
# disconnect
player.channel_id = None
self.cleanup()