Implemented lavalink for bot + cog

This commit is contained in:
2021-10-31 16:32:56 +01:00
parent fc7da80b29
commit 75feac1e0e
2 changed files with 147 additions and 70 deletions
+4 -2
View File
@@ -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", ""),
+141 -66
View File
@@ -1,45 +1,91 @@
import discord
from discord import channel
from discord.ext import commands
import re
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"):
bot.lavalink = lavalink.Client(bot.user.id)
bot.lavalink.add_node("de-1.rivalmc.net", 2333, "12345", "eu", "poggers")
bot.add_listener(
self.bot.lavalink.voice_update_handler, "on_socket_response"
@commands.Cog.listener()
async def on_ready(self):
if not hasattr(
self.bot, "lavalink"
): # 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)
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])
self.bot.lavalink.add_event_hook(self.track_hook)
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):
"""Connects to the given voicechannel ID. A channel_id of `None` means disconnect."""
ws = self.bot._connection._get_websocket(guild_id)
await ws.voice_state(str(guild_id), channel_id)
@commands.command(name="play", aliases=["p", "connect", "join"])
@commands.guild_only()
async def play(self, ctx: Context):
"""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")
async def connect(self, ctx):
"""Starts vibing."""
player = self.bot.lavalink.player_manager.get(ctx.guild.id)
results = await player.node.get_tracks(self.stream)
# Get the results for the query from Lavalink.
results = await player.node.get_tracks(query)
# 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":
tracks = results["tracks"]
for track in tracks:
player.add(requester=ctx.author.id, track=track)
else:
# Theoretically songs will always be TRACK_LOADED
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"])
async def disconnect(self, ctx):
@commands.command(name="disconnect", aliases=["dc", "stop"])
@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)
@commands.command(name="now", aliases=["playing"])
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)
# Disconnect from the voice channel.
await ctx.voice_client.disconnect(force=True)
await EmbedGenerator.Title(ctx, "*⃣ | Disconnected.")
def setup(bot: ChristmasBot):
bot.add_cog(MusicCog(bot))
bot.add_cog(Music(bot))