musicube/src/bot.py

import os
import sys
import asyncio

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

from error import ErrorHandler

load_dotenv()

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


class Music(commands.Cog):
  def __init__(self, bot):
    self._bot = bot
    self._queue = asyncio.Queue()
    self._queue_lock = asyncio.Lock()
    self._ytdl_format_options = {
        '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_format_options)
    # pylint: disable=no-member
    self._handle_playback.start()

  @tasks.loop()
  async def _handle_playback(self):
    while True:
      await self._queue_lock.acquire()
      ctx, url, title = await self._queue.get()
      audio = discord.FFmpegPCMAudio(url, **self._ffmpeg_options)

      def after(err):
        if err:
          print(f"Player error: {err}")
        self._queue.task_done()
        self._queue_lock.release()
      ctx.voice_client.play(audio, after=after)
      await ctx.send(f"Now playing: {title}")

  @commands.command()
  async def play(self, ctx, *, url):
    async with ctx.typing():
      data = self._ytdl.extract_info(url, download=False)
      if 'entries' in data:
        data = data['entries'][0]
        title = data.get('title')
        url = data.get('url')
        await self._queue.put((ctx, url, title))
        if ctx.voice_client.is_playing():
          await ctx.send(f"Queued: {title}")

  @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("BOT_TOKEN", None)
  if not token:
    print("No token found in BOT_TOKEN")
    sys.exit(1)

  bot.run(token)