diff --git a/cogs/information.py b/cogs/information.py index 3cebf27..4cf85fc 100644 --- a/cogs/information.py +++ b/cogs/information.py @@ -9,6 +9,7 @@ import humanize import datetime import lavalink from bot import ChristmasBot +from context import CustomContext from utils.EmbedGenerator import EmbedGenerator from utils.paginator import HelpPaginator from utils.database import AutoJoin @@ -56,9 +57,9 @@ class InformationCog(commands.Cog, name="Information"): slash_commands=True, ) @commands.cooldown(rate=1, per=5, type=commands.BucketType.user) - async def wlinfo(self, ctx: Context): + async def wlinfo(self, ctx: CustomContext): """Retrieve various Node/Server/Player information.""" - player = self.bot.lavalink.player_manager.get(ctx.guild.id) + player = ctx.get_player() node = player.node used = humanize.naturalsize(node.stats.memory_used) diff --git a/cogs/music.py b/cogs/music.py index e577a34..2abe7bf 100644 --- a/cogs/music.py +++ b/cogs/music.py @@ -217,7 +217,7 @@ class Music(commands.Cog): async def play(self, ctx: CustomContext): """Starts playing Christmas bangers""" # Get the player for this guild from cache. - player: DefaultPlayer = self.bot.lavalink.player_manager.get(ctx.guild.id) + player = ctx.get_player() await self.fill_player_queue(player, self.bot.config["queue_buffer_size"]) if not player.is_playing: @@ -226,22 +226,22 @@ class Music(commands.Cog): await ctx.send("Started playing") @commands.command(name="skip", aliases=["next"]) - async def skip(self, ctx: Context): + async def skip(self, ctx: CustomContext): """I heard this song way too often""" - player: DefaultPlayer = self.bot.lavalink.player_manager.get(ctx.guild.id) + player = ctx.get_player() await player.skip() await ctx.send("Skipped current song") @commands.command(name="queue") - async def queue(self, ctx: Context): + async def queue(self, ctx: CustomContext): """Ghetto queue""" - player: DefaultPlayer = self.bot.lavalink.player_manager.get(ctx.guild.id) + player = ctx.get_player() await EmbedGenerator.Message(ctx, "Queue:", player.queue) @commands.command(name="disconnect", aliases=["dc", "stop"]) - async def disconnect(self, ctx: Context): + async def disconnect(self, ctx: CustomContext): """Disconnects ChristmasBot""" - player: DefaultPlayer = self.bot.lavalink.player_manager.get(ctx.guild.id) + player = ctx.get_player() if not player.is_connected: return await EmbedGenerator.Title(ctx, "Not connected.") diff --git a/context.py b/context.py index c518260..d5051ea 100644 --- a/context.py +++ b/context.py @@ -1,7 +1,15 @@ from aioredis.client import Redis from discord.ext import commands +from discord.ext.commands.errors import CommandInvokeError +from lavalink.models import DefaultPlayer class CustomContext(commands.Context): def get_redis(self) -> Redis: return self.bot._redis_client + + def get_player(self) -> DefaultPlayer: + if hasattr(self.bot, "lavalink"): + return self.bot.lavalink.player_manager.get(self.guild.id) + + raise CommandInvokeError("Lavalink is still starting up.")