Implemented lavalink for bot + cog

bot.py

@@ -11,9 +11,13 @@ from discord.ext.commands.errors import (
     ExtensionNotFound,
     NoEntryPointError,
 )
+import lavalink
 
 
 class ChristmasBot(commands.Bot):
+    lavalink: lavalink.Client
+    invite_link: str
+
     def __init__(self, config: Dict[Any, Any]):
         intents = discord.Intents(
             voice_states=True, guild_messages=True, guilds=True, messages=True
@@ -23,8 +27,6 @@ class ChristmasBot(commands.Bot):
         self.initial_cog_names: List[str] = self.config.get("cogs", [])
         self.colors: Dict[str, Color] = self.process_colours(config.get("colors", []))
 
-        self.invite_link: str = ""
-
         super().__init__(
             command_prefix=self.prefix_callable,
             description=self.config["info"].get("description", ""),

cogs/music.py

@@ -1,45 +1,91 @@
-import discord
+import re
-from discord import channel
-from discord.ext import commands
 
+import discord
+from discord.ext.commands.context import Context
 import lavalink
-import asyncio
+from discord.ext import commands
+from lavalink.models import DefaultPlayer
+
 from bot import ChristmasBot
 from utils.EmbedGenerator import EmbedGenerator
 
+url_rx = re.compile(r"https?://(?:www\.)?.+")
 
-class MusicCog(commands.Cog, name="Music"):
+
-    def __init__(self, bot: commands.Bot):
+class LavalinkVoiceClient(discord.VoiceClient):
+    def __init__(self, client: discord.Client, channel: discord.abc.Connectable):
+        self.client = client
+        self.channel = channel
+        # ensure there exists a client already
+        if hasattr(self.client, "lavalink"):
+            self.lavalink = self.client.lavalink
+        else:
+            self.client.lavalink = lavalink.Client(client.user.id)
+            self.client.lavalink.add_node(
+                "localhost", 2333, "youshallnotpass", "us", "default-node"
+            )
+            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) -> 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)
+
+    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()
+
+
+class Music(commands.Cog):
+    def __init__(self, bot: ChristmasBot):
         self.bot = bot
-        self.stream = "https://azuracast.exobot.site/radio/8000/radio.opus"
 
-        if not hasattr(bot, "lavalink"):
+    @commands.Cog.listener()
-            bot.lavalink = lavalink.Client(bot.user.id)
+    async def on_ready(self):
-            bot.lavalink.add_node("de-1.rivalmc.net", 2333, "12345", "eu", "poggers")
+        if not hasattr(
-            bot.add_listener(
+            self.bot, "lavalink"
-                self.bot.lavalink.voice_update_handler, "on_socket_response"
+        ):  # This ensures the client isn't overwritten during cog reloads.
+            self.bot.lavalink = lavalink.Client(self.bot.user.id)
+
+            ll = self.bot.config["lavalink"]
+            self.bot.lavalink.add_node(
+                ll["host"], ll["port"], ll["password"], ll["region"], ll["name"]
             )
 
-        lavalink.add_event_hook(self.track_hook)
+        self.bot.lavalink.add_event_hook(self.track_hook)
-
-        bot.loop.create_task(self.async_init())
-
-    async def async_init(self):
-        await self.bot.wait_until_ready()
-
-        channels = []
-
-        # We startup to fast #NOTPOGGERS
-        await asyncio.sleep(5)
-        for channel in channels:
-            guild = self.bot.get_guild(channel[0])
-            player = self.bot.lavalink.player_manager.create(
-                channel[0], endpoint=str(guild.region)
-            )
-            track = await player.node.get_tracks(self.stream)
-            if not player.is_playing:
-                await player.play(track["tracks"][0])
-            await self.connect_to(channel[0], channel[1])
 
     def cog_unload(self):
         """Cog unload handler. This removes any event hooks that were registered."""
@@ -48,18 +94,42 @@ class MusicCog(commands.Cog, name="Music"):
     async def cog_before_invoke(self, ctx):
         """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, error):
+        if isinstance(error, commands.CommandInvokeError):
+            await ctx.send(error.original)
+            # The above handles errors thrown in this cog and shows them to the user.
+            # This shouldn't be a problem as the only errors thrown in this cog are from `ensure_voice`
+            # which contain a reason string, such as "Join a voicechannel" etc. You can modify the above
+            # if you want to do things differently.
+
     async def ensure_voice(self, ctx):
         """This check ensures that the bot and command author are in the same voicechannel."""
         player = self.bot.lavalink.player_manager.create(
             ctx.guild.id, endpoint=str(ctx.guild.region)
         )
-        should_connect = ctx.command.name in ("connect",)
+        # 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 ("play",)
 
         if not ctx.author.voice or not ctx.author.voice.channel:
+            # Our cog_command_error handler catches this and sends it to the voicechannel.
+            # Exceptions allow us to "short-circuit" command invocation via checks so the
+            # execution state of the command goes no further.
             raise commands.CommandInvokeError("Join a voicechannel first.")
 
         if not player.is_connected:
@@ -76,69 +146,74 @@ class MusicCog(commands.Cog, name="Music"):
                 )
 
             player.store("channel", ctx.channel.id)
-            await self.connect_to(ctx.guild.id, str(ctx.author.voice.channel.id))
+            await ctx.author.voice.channel.connect(cls=LavalinkVoiceClient)
         else:
             if int(player.channel_id) != ctx.author.voice.channel.id:
                 raise commands.CommandInvokeError("You need to be in my voicechannel.")
 
     async def track_hook(self, 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 = int(event.player.guild_id)
-            await self.connect_to(guild_id, None)
+            guild = self.bot.get_guild(guild_id)
+            await guild.voice_client.disconnect(force=True)
 
-    async def connect_to(self, guild_id: int, channel_id: str):
+    @commands.command(name="play", aliases=["p", "connect", "join"])
-        """Connects to the given voicechannel ID. A channel_id of `None` means disconnect."""
+    @commands.guild_only()
-        ws = self.bot._connection._get_websocket(guild_id)
+    async def play(self, ctx: Context):
-        await ws.voice_state(str(guild_id), channel_id)
+        """Searches and plays a song from a given query."""
+        # Get the player for this guild from cache.
+        player: DefaultPlayer = self.bot.lavalink.player_manager.get(ctx.guild.id)
+        query = "https://youtu.be/VcIt_AcOPjs"
 
-    @commands.command(name="connect")
+        # Get the results for the query from Lavalink.
-    async def connect(self, ctx):
+        results = await player.node.get_tracks(query)
-        """Starts vibing."""
-        player = self.bot.lavalink.player_manager.get(ctx.guild.id)
-        results = await player.node.get_tracks(self.stream)
 
+        # Results could be None if Lavalink returns an invalid response (non-JSON/non-200 (OK)).
+        # ALternatively, resullts['tracks'] could be an empty array if the query yielded no tracks.
         if not results or not results["tracks"]:
-            return await ctx.send("Nothing found!")
+            return await EmbedGenerator.Title(ctx, "Nothing found!")
 
-        if results["loadType"] == "PLAYLIST_LOADED":
+        # Theoretically songs will always be TRACK_LOADED
-            tracks = results["tracks"]
-
-            for track in tracks:
-                player.add(requester=ctx.author.id, track=track)
-        else:
         track = results["tracks"][0]
+        await EmbedGenerator.Message(
+            ctx, "Track Enqueued", f'[{track["info"]["title"]}]({track["info"]["uri"]})'
+        )
+
         track = lavalink.models.AudioTrack(track, ctx.author.id, recommended=True)
         player.add(requester=ctx.author.id, track=track)
 
         if not player.is_playing:
             await player.play()
 
-    @commands.command(aliases=["dc"])
+    @commands.command(name="disconnect", aliases=["dc", "stop"])
-    async def disconnect(self, ctx):
+    @commands.guild_only()
+    async def disconnect(self, ctx: Context):
         """Disconnects the player from the voice channel and clears its queue."""
-        player = self.bot.lavalink.player_manager.get(ctx.guild.id)
+        player: DefaultPlayer = self.bot.lavalink.player_manager.get(ctx.guild.id)
 
         if not player.is_connected:
-            return await ctx.send("Not connected.")
+            return await EmbedGenerator.Title(ctx, "Not connected.")
 
         if not ctx.author.voice or (
             player.is_connected
             and ctx.author.voice.channel.id != int(player.channel_id)
         ):
-            return await ctx.send("You're not in my voicechannel!")
+            # Abuse prevention. Users not in voice channels, or not in the same voice channel as the bot
+            # may not disconnect the bot.
+            return await EmbedGenerator.Title(ctx, "You're not in my voicechannel!")
 
+        # 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()
-        await self.connect_to(ctx.guild.id, None)
+        # Disconnect from the voice channel.
-
+        await ctx.voice_client.disconnect(force=True)
-    @commands.command(name="now", aliases=["playing"])
+        await EmbedGenerator.Title(ctx, "*⃣ | Disconnected.")
-    async def now_playing(self, ctx):
-        """Stop and disconnect the player and controller."""
-        em = discord.Embed(color=self.bot.colors["embed"])
-        # em.set_thumbnail(url=np["thumbnail"])
-        em.add_field(name="Currently playing:", value="Some song")
-        await EmbedGenerator.SendWithFooter(ctx, em)
 
 
 def setup(bot: ChristmasBot):
-    bot.add_cog(MusicCog(bot))
+    bot.add_cog(Music(bot))