musicube/src/bot.py

import os
import sys
import asyncio
from dataclasses import dataclass

import discord
from discord.ext import commands, tasks
from dotenv import load_dotenv
import youtube_dl

from error import ErrorHandler
from message import NowPlayingMessage, QueuedMessage, ErrorMessage

load_dotenv()

# Suppress noise about console usage from errors
youtube_dl.utils.bug_reports_message = lambda: ''


@dataclass
class Song:
  title: str
  webpage_url: str
  audio_url: str


class Music(commands.Cog):
  def __init__(self, bot):
    self._bot = bot
    self._queue = asyncio.Queue()
    self._queue_lock = asyncio.Lock()
    self._current_skip_message = None
    self._ytdl_params = {
        'cookiefile': os.environ.get('YOUTUBE_COOKIES'),
        'format': 'bestaudio/best',
        'outtmpl': '%(extractor)s-%(id)s-%(title)s.%(ext)s',
        'restrictfilenames': True,
        'noplaylist': True,
        'nocheckcertificate': True,
        'ignoreerrors': False,
        'logtostderr': False,
        'quiet': True,
        'no_warnings': True,
        'default_search': 'auto',
        'source_address': '0.0.0.0'  # bind to ipv4 since ipv6 addresses cause issues sometimes
    }
    self._ffmpeg_options = {
        'options': '-vn'
    }
    self._ytdl = youtube_dl.YoutubeDL(self._ytdl_params)
    # pylint: disable=no-member
    self._handle_playback.start()

  def cog_unload(self):
    # pylint: disable=no-member
    self._handle_playback.cancel()

  def _next(self):
    """Trigger playback of next song."""
    self._queue.task_done()
    self._current_skip_message = None
    if self._queue_lock.locked():
      self._queue_lock.release()

  @tasks.loop()
  async def _handle_playback(self):
    while True:
      try:
        await self._queue_lock.acquire()
        ctx, song = await self._queue.get()
        if ctx.voice_client is None:
          # Bot is no longer in a voice channel.
          # This could be the case because a stop command was issued.
          # We will skip this (and possibly all remaining songs) in the queue
          self._next()
          continue
        audio = discord.FFmpegPCMAudio(song.audio_url, **self._ffmpeg_options)

        def after(err):
          if err:
            print(f"Player error: {err}")
          self._next()
        ctx.voice_client.play(audio, after=after)
        embed = NowPlayingMessage(title=song.title, url=song.webpage_url)
        msg = await ctx.send(embed=embed)
        await self._add_skip_button(msg)
      # pylint: disable=broad-except
      except Exception as err:
        print(f"Error during playback: {err}")
        if ctx:
          embed = ErrorMessage(str(err))
          await ctx.send(embed=embed)
        self._next()

  async def _add_skip_button(self, msg):
    await msg.add_reaction("⏭️")
    self._current_skip_message = msg

  @commands.Cog.listener()
  async def on_reaction_add(self, reaction, user):
    if not user.bot \
            and self._current_skip_message \
            and reaction.message.id == self._current_skip_message.id \
            and reaction.emoji == "⏭️":
      voice_client = reaction.message.guild.voice_client
      self._skip(voice_client)

  @_handle_playback.before_loop
  async def before_handle_playback(self):
    await self._bot.wait_until_ready()

  @commands.command()
  async def play(self, ctx, *, query):
    async with ctx.typing():
      data = self._ytdl.extract_info(query, download=False)
      if 'entries' in data:
        data = data['entries'][0]
        song = Song(
            title=data.get('title'),
            audio_url=data.get('url'),
            webpage_url=data.get('webpage_url')
        )
        await self._queue.put((ctx, song))
        if ctx.voice_client.is_playing():
          embed = QueuedMessage(title=song.title, url=song.webpage_url)
          await ctx.send(embed=embed)

  def _skip(self, voice_client):
    """Skip to next song."""
    if voice_client is None or not voice_client.is_playing():
      raise commands.CommandError("No song playing")
    # This skips to next song because the bot does not differentiate between
    # a song stopping because it is finished or because it was manually stopped.
    voice_client.stop()

  @commands.command()
  async def skip(self, ctx):
    async with ctx.typing():
      self._skip(ctx.voice_client)

  @commands.command()
  async def stop(self, ctx):
    await ctx.voice_client.disconnect()

  @play.before_invoke
  async def ensure_voice(self, ctx):
    if ctx.voice_client is None:
      if ctx.author.voice:
        await ctx.author.voice.channel.connect()
      else:
        raise commands.CommandError("Author not connected to a voice channel")


if __name__ == "__main__":
  bot = commands.Bot(command_prefix=commands.when_mentioned_or("!"), description='Relatively simple music bot example')

  @bot.event
  async def on_ready():
    print(f"Logged in as {bot.user} ({bot.user.id})")
    print('------')

  bot.add_cog(Music(bot))
  bot.add_cog(ErrorHandler(bot))

  token = os.environ.get("DISCORD_BOT_TOKEN", None)
  if not token:
    print("Discord bot token not found")
    sys.exit(1)

  bot.run(token)