From 316f3bacc92991a70329a6c08cf1d23c2a401f91 Mon Sep 17 00:00:00 2001 From: strNophix Date: Thu, 18 Nov 2021 21:59:37 +0100 Subject: [PATCH] Added sources functionality + pre-commit --- .pre-commit-config.yaml | 2 +- cogs/information.py | 2 +- cogs/music.py | 8 +-- cogs/settings.py | 116 +++++++++++++++++++++++++++++++++-- context.py | 18 ++++++ scripts/channel_to_redis.py | 6 +- scripts/playlist_to_redis.py | 7 ++- utils/EmbedGenerator.py | 5 +- utils/classes.py | 2 + utils/database.py | 22 +++++++ utils/exceptions.py | 3 +- 11 files changed, 170 insertions(+), 21 deletions(-) diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index 76a1554..7e3b957 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -1,5 +1,5 @@ default_language_version: - python: python3.8 + python: python3.9 repos: - repo: https://github.com/pre-commit/pre-commit-hooks rev: v4.0.1 diff --git a/cogs/information.py b/cogs/information.py index bbdab08..a8ba045 100644 --- a/cogs/information.py +++ b/cogs/information.py @@ -11,9 +11,9 @@ from discord.ext.commands import Context from bot import TuneBot from context import CustomContext +from utils.classes import BaseCog from utils.database import AutoJoin from utils.EmbedGenerator import EmbedGenerator -from utils.classes import BaseCog from utils.paginator import HelpPaginator diff --git a/cogs/music.py b/cogs/music.py index 02cb4ee..282cb67 100644 --- a/cogs/music.py +++ b/cogs/music.py @@ -4,24 +4,22 @@ import re from typing import Optional import discord -from discord.channel import TextChannel -from discord.ext.commands.context import Context -from discord.ext.commands.errors import CommandError import lavalink from discord import Embed from discord.channel import TextChannel from discord.ext import commands from discord.ext.commands.context import Context +from discord.ext.commands.errors import CommandError from lavalink.models import AudioTrack from lavalink.models import DefaultPlayer from bot import TuneBot -from utils.classes import BaseCog from context import CustomContext +from utils.classes import BaseCog from utils.database import AutoJoin from utils.database import Playlist -from utils.exceptions import EmbeddedCommandException from utils.EmbedGenerator import EmbedGenerator +from utils.exceptions import EmbeddedCommandException url_rx = re.compile(r"https?://(?:www\.)?.+") diff --git a/cogs/settings.py b/cogs/settings.py index fee02a4..725cf53 100644 --- a/cogs/settings.py +++ b/cogs/settings.py @@ -1,13 +1,14 @@ +from typing import Any + from discord.ext import commands -from utils.EmbedGenerator import EmbedGenerator -from utils.classes import BaseCog -from utils.database import AutoJoin -from bot import TuneBot -from discord.ext.commands import Context +from discord.message import Message from bot import TuneBot from context import CustomContext +from utils.classes import BaseCog from utils.database import AutoJoin +from utils.database import Playlist +from utils.database import PlaylistSource from utils.EmbedGenerator import EmbedGenerator @@ -42,6 +43,111 @@ class SettingsCog(BaseCog, name="Settings"): await AutoJoin.del_channel(ctx.redis, ctx.guild.id) await EmbedGenerator.Message(ctx, "Autojoin", "`disabled`") + @commands.group( + name="source", + aliases=["src"], + invoke_without_command=True, + slash_command=False, + ) + async def source(self, ctx: CustomContext): + """Displays all possible options for the `source` command""" + prefix = self.bot.config["prefixes"][0] + embed = ctx.create_embed() + embed.title = "All options:" + embed.description = f"```{prefix}source list\n{prefix}source add \n{prefix}source remove \n{prefix}source sync```" + await ctx.send(embed=embed) + + @source.command(name="remove") + async def source_remove(self, ctx: CustomContext, source_url: str): + """Removes a source from the bot""" + if await PlaylistSource.remove(ctx.redis, source_url): + prefix = self.bot.config["prefixes"][0] + embed = ctx.create_embed() + embed.title = "Removed source succesfully" + embed.description = ( + f"Please use `{prefix}source sync` to persist these changes." + ) + await ctx.send(embed=embed) + return + + embed = ctx.create_embed() + embed.title = "Could not remove source, the specified source might not exist" + await ctx.send(embed=embed) + + @source.command(name="add") + async def source_add(self, ctx: CustomContext, source_url: str): + """ + Add's a source to the bot + + Supported sources: YouTube, SoundCloud, Bandcamp, Vimeo, Twitch and HTTP(S) URL's + """ + embed = ctx.create_embed() + embed.title = "Started processing source" + message = await ctx.send(embed=embed) + + query_result: Any = await self.bot.lavalink.get_tracks(source_url) + + if query_result["loadType"] == "LOAD_FAILED": + embed = ctx.create_embed() + embed.title = "The specified URL is not a valid source" + embed.description = f"Supported sources: YouTube, SoundCloud, Bandcamp, Vimeo, Twitch and HTTP(S) URL's" + if isinstance(message, Message): + await message.edit(embed=embed) + else: + await ctx.send(embed=embed) + return + + await PlaylistSource.add(ctx.redis, source_url) + track_urls = [str(track["info"]["uri"]) for track in query_result["tracks"]] + await Playlist.add_bulk(ctx.redis, track_urls) + + embed = ctx.create_embed() + embed.title = "Finished processing source" + embed.description = f"Added {len(track_urls)} tracks" + if isinstance(message, Message): + await message.edit(embed=embed) + else: + await ctx.send(embed=embed) + + @source.command(name="list", aliases=["ls"]) + async def source_list(self, ctx: CustomContext): + """Display a list of sources""" + # TODO: Implement pagination for sources + sources = await PlaylistSource.get_all(ctx.redis) + if len(sources) > 0: + description = "\n".join([f"[{source}]({source})" for source in sources]) + else: + description = "This bot has no sources yet" + + embed = ctx.create_embed() + embed.title = "All sources:" + embed.description = description + await ctx.send(embed=embed) + + @source.command(name="sync") + async def source_sync(self, ctx: CustomContext): + """Forcefully resyncs all sources""" + failed_sources: list[str] = [] + await Playlist.clear(ctx.redis) + sources = await PlaylistSource.get_all(ctx.redis) + for source_url in sources: + query_result: Any = await self.bot.lavalink.get_tracks(source_url) + if query_result["loadType"] == "LOAD_FAILED": + failed_sources.append(source_url) + continue + + track_urls = [str(track["info"]["uri"]) for track in query_result["tracks"]] + await Playlist.add_bulk(ctx.redis, track_urls) + + embed = ctx.create_embed() + embed.title = f"Finished sync ({len(failed_sources)} issues)" + if len(failed_sources) > 0: + embed.description = "\n".join( + [f"[{source_url}]({source_url})" for source_url in failed_sources] + ) + + await ctx.send(embed=embed) + def setup(bot: TuneBot): bot.add_cog(SettingsCog(bot)) diff --git a/context.py b/context.py index 3f41dc7..433fc9e 100644 --- a/context.py +++ b/context.py @@ -1,8 +1,14 @@ +from typing import TYPE_CHECKING + from aioredis.client import Redis +from discord import Embed from discord.ext import commands from discord.ext.commands.errors import CommandInvokeError from lavalink.models import DefaultPlayer +if TYPE_CHECKING: + from bot import TuneBot + class CustomContext(commands.Context): @property @@ -15,3 +21,15 @@ class CustomContext(commands.Context): return self.bot.lavalink.player_manager.get(self.guild.id) raise CommandInvokeError("Lavalink is still starting up.") + + def create_embed(self) -> Embed: + bot: TuneBot = self.bot + color = bot.colors["embed"] + + avatar = None + if avatar_asset := self.author.avatar: + avatar = avatar_asset.with_static_format("jpeg") + + embed = Embed(color=color) + embed.set_footer(text=f"Requested by: {self.author}", icon_url=avatar) + return embed diff --git a/scripts/channel_to_redis.py b/scripts/channel_to_redis.py index 4f10e25..4c4315e 100644 --- a/scripts/channel_to_redis.py +++ b/scripts/channel_to_redis.py @@ -1,8 +1,9 @@ -import subprocess import json +import subprocess +import sys + import redis from yt_dlp import YoutubeDL -import sys if len(sys.argv) < 2: raise Exception("Expected youtube playlist/channel/video") @@ -25,4 +26,3 @@ for vid_url in vid_urls.split("\n"): vid_url = "https://www.youtube.com/watch?v=" + vid_url redis_client.sadd(f"{redis_prefix}:playlist", vid_url) print(vid_url) - diff --git a/scripts/playlist_to_redis.py b/scripts/playlist_to_redis.py index b5fecd7..be97a7a 100644 --- a/scripts/playlist_to_redis.py +++ b/scripts/playlist_to_redis.py @@ -1,8 +1,9 @@ import json -from aiotube import Playlist -import redis import sys +import redis +from aiotube import Playlist + if len(sys.argv) < 2: raise Exception("Expected path to file as argument") @@ -18,4 +19,4 @@ for line in file: for vid in playlist.videos(): yt_url = vid.url redis_client.sadd(f"{redis_prefix}:playlist", yt_url) - print(yt_url) \ No newline at end of file + print(yt_url) diff --git a/utils/EmbedGenerator.py b/utils/EmbedGenerator.py index 95cf8c0..1ac67e5 100644 --- a/utils/EmbedGenerator.py +++ b/utils/EmbedGenerator.py @@ -1,8 +1,9 @@ -from typing import Optional, Union +from typing import Optional +from typing import Union import discord -from discord.ext.commands import Context from discord import Embed +from discord.ext.commands import Context class EmbedGenerator: diff --git a/utils/classes.py b/utils/classes.py index 5c90b02..9cc95b2 100644 --- a/utils/classes.py +++ b/utils/classes.py @@ -1,5 +1,7 @@ from typing import Dict + from discord.ext.commands import Cog + from bot import TuneBot diff --git a/utils/database.py b/utils/database.py index 0d48ec6..b8bcc06 100644 --- a/utils/database.py +++ b/utils/database.py @@ -33,3 +33,25 @@ class Playlist: @staticmethod async def random(redis: Redis, amount: Optional[int] = 1) -> List[str]: return await redis.srandmember(f"{redis_prefix}:playlist", amount) + + @staticmethod + async def add_bulk(redis: Redis, urls: list[str]): + await redis.sadd(f"{redis_prefix}:playlist", *urls) + + @staticmethod + async def clear(redis: Redis): + await redis.delete(f"{redis_prefix}:playlist") + + +class PlaylistSource: + @staticmethod + async def get_all(redis: Redis) -> list[str]: + return await redis.smembers(f"{redis_prefix}:sources") + + @staticmethod + async def add(redis: Redis, source_url: str): + await redis.sadd(f"{redis_prefix}:sources", source_url) + + @staticmethod + async def remove(redis: Redis, source_url: str) -> bool: + return await redis.srem(f"{redis_prefix}:sources", source_url) diff --git a/utils/exceptions.py b/utils/exceptions.py index d807e17..b6a08e0 100644 --- a/utils/exceptions.py +++ b/utils/exceptions.py @@ -1,7 +1,8 @@ from discord.embeds import Embed -from context import CustomContext from discord.ext.commands import CommandError +from context import CustomContext + class EmbeddedCommandException(CommandError): def __init__(self, embed: Embed) -> None: