Merge pull request #35 from strNophix/source-management

Source management

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
 
 

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\.)?.+")
 

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,117 @@ class SettingsCog(BaseCog, name="Settings"):
         await AutoJoin.del_channel(ctx.redis, ctx.guild.id)
         await EmbedGenerator.Message(ctx, "Autojoin", "`disabled`")
 
+    @commands.is_owner()
+    @commands.group(
+        name="source",
+        aliases=["src"],
+        invoke_without_command=True,
+        slash_command=False,
+        hidden=True,
+    )
+    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 <url>\n{prefix}source remove <url>\n{prefix}source sync```"
+        await ctx.send(embed=embed)
+
+    @commands.is_owner()
+    @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)
+
+    @commands.is_owner()
+    @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)
+
+    @commands.is_owner()
+    @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)
+
+    @commands.is_owner()
+    @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))

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

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)
-

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")
 

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:

utils/classes.py

@@ -1,5 +1,7 @@
 from typing import Dict
+
 from discord.ext.commands import Cog
+
 from bot import TuneBot
 
 

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)

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: