From 1ab823569e4d0f86bed7ba9020485566ba9ba45f Mon Sep 17 00:00:00 2001 From: Lemon4ksan Date: Mon, 3 Feb 2025 22:31:06 +0300 Subject: [PATCH] impr: Async database and code optimization. --- MusicBot/cogs/general.py | 18 +- MusicBot/cogs/settings.py | 28 ++- MusicBot/cogs/utils/embeds.py | 4 +- MusicBot/cogs/utils/voice_extension.py | 236 ++++++++++-------- MusicBot/cogs/voice.py | 60 +++-- MusicBot/database/__init__.py | 6 +- MusicBot/database/base.py | 272 ++++++++------------ MusicBot/database/extensions.py | 329 +++++++++++++++---------- MusicBot/ui/find.py | 46 ++-- MusicBot/ui/menu.py | 44 ++-- MusicBot/ui/other.py | 91 ++++--- 11 files changed, 600 insertions(+), 534 deletions(-) diff --git a/MusicBot/cogs/general.py b/MusicBot/cogs/general.py index ad38cdb..3d5d873 100644 --- a/MusicBot/cogs/general.py +++ b/MusicBot/cogs/general.py @@ -23,7 +23,7 @@ async def get_search_suggestions(ctx: discord.AutocompleteContext) -> list[str]: return [] users_db = BaseUsersDatabase() - token = users_db.get_ym_token(ctx.interaction.user.id) + token = await users_db.get_ym_token(ctx.interaction.user.id) if not token: logging.info(f"[GENERAL] User {ctx.interaction.user.id} has no token") return [] @@ -186,21 +186,21 @@ class General(Cog): about = cast(yandex_music.Status, client.me).to_dict() uid = ctx.author.id - self.users_db.update(uid, {'ym_token': token}) + await self.users_db.update(uid, {'ym_token': token}) logging.info(f"[GENERAL] Token saved for user {ctx.author.id}") await ctx.respond(f'Привет, {about['account']['first_name']}!', delete_after=15, ephemeral=True) @account.command(description="Удалить токен из датабазы бота.") async def remove(self, ctx: discord.ApplicationContext) -> None: logging.info(f"[GENERAL] Remove command invoked by user {ctx.author.id} in guild {ctx.guild.id}") - self.users_db.update(ctx.user.id, {'ym_token': None}) + await self.users_db.update(ctx.user.id, {'ym_token': None}) await ctx.respond(f'Токен был удалён.', delete_after=15, ephemeral=True) @account.command(description="Получить плейлист «Мне нравится»") async def likes(self, ctx: discord.ApplicationContext) -> None: logging.info(f"[GENERAL] Likes command invoked by user {ctx.author.id} in guild {ctx.guild.id}") - token = self.users_db.get_ym_token(ctx.user.id) + token = await self.users_db.get_ym_token(ctx.user.id) if not token: logging.info(f"[GENERAL] No token found for user {ctx.user.id}") await ctx.respond("❌ Укажите токен через /account login.", delete_after=15, ephemeral=True) @@ -232,7 +232,7 @@ class General(Cog): async def playlists(self, ctx: discord.ApplicationContext) -> None: logging.info(f"[GENERAL] Playlists command invoked by user {ctx.user.id} in guild {ctx.guild_id}") - token = self.users_db.get_ym_token(ctx.user.id) + token = await self.users_db.get_ym_token(ctx.user.id) if not token: await ctx.respond("❌ Укажите токен через /account login.", delete_after=15, ephemeral=True) return @@ -247,11 +247,11 @@ class General(Cog): (playlist.title if playlist.title else 'Без названия', playlist.track_count if playlist.track_count else 0) for playlist in playlists_list ] - self.users_db.update(ctx.user.id, {'playlists': playlists, 'playlists_page': 0}) + await self.users_db.update(ctx.user.id, {'playlists': playlists, 'playlists_page': 0}) embed = generate_playlists_embed(0, playlists) logging.info(f"[GENERAL] Successfully fetched playlists for user {ctx.user.id}") - await ctx.respond(embed=embed, view=MyPlaylists(ctx), ephemeral=True) + await ctx.respond(embed=embed, view=await MyPlaylists(ctx).init(), ephemeral=True) @discord.slash_command(description="Найти контент и отправить информацию о нём. Возвращается лучшее совпадение.") @discord.option( @@ -276,8 +276,8 @@ class General(Cog): ) -> None: logging.info(f"[GENERAL] Find command invoked by user {ctx.user.id} in guild {ctx.guild_id} for '{content_type}' with name '{name}'") - guild = self.db.get_guild(ctx.guild_id) - token = self.users_db.get_ym_token(ctx.user.id) + guild = await self.db.get_guild(ctx.guild_id, projection={'allow_explicit': 1}) + token = await self.users_db.get_ym_token(ctx.user.id) if not token: logging.info(f"[GENERAL] No token found for user {ctx.user.id}") await ctx.respond("❌ Укажите токен через /account login.", delete_after=15, ephemeral=True) diff --git a/MusicBot/cogs/settings.py b/MusicBot/cogs/settings.py index 92ae96a..57c0e11 100644 --- a/MusicBot/cogs/settings.py +++ b/MusicBot/cogs/settings.py @@ -19,7 +19,9 @@ class Settings(Cog): @settings.command(name="show", description="Показать текущие настройки бота.") async def show(self, ctx: discord.ApplicationContext) -> None: - guild = self.db.get_guild(ctx.guild.id) + guild = await self.db.get_guild(ctx.guild.id, projection={ + 'allow_explicit': 1, 'always_allow_menu': 1, 'vote_next_track': 1, 'vote_add_track': 1, 'vote_add_album': 1, 'vote_add_artist': 1, 'vote_add_playlist': 1 + }) embed = discord.Embed(title="Настройки бота", color=0xfed42b) explicit = "✅ - Разрешены" if guild['allow_explicit'] else "❌ - Запрещены" @@ -44,8 +46,8 @@ class Settings(Cog): await ctx.respond("❌ У вас нет прав для выполнения этой команды.", delete_after=15, ephemeral=True) return - guild = self.db.get_guild(ctx.guild.id) - self.db.update(ctx.guild.id, {'allow_explicit': not guild['allow_explicit']}) + guild = await self.db.get_guild(ctx.guild.id, projection={'allow_explicit': 1}) + await self.db.update(ctx.guild.id, {'allow_explicit': not guild['allow_explicit']}) await ctx.respond(f"Треки с содержанием не для детей теперь {'разрешены' if not guild['allow_explicit'] else 'запрещены'}.", delete_after=15, ephemeral=True) @settings.command(name="menu", description="Разрешить или запретить создание меню проигрывателя, даже если в канале больше одного человека.") @@ -55,8 +57,8 @@ class Settings(Cog): await ctx.respond("❌ У вас нет прав для выполнения этой команды.", delete_after=15, ephemeral=True) return - guild = self.db.get_guild(ctx.guild.id) - self.db.update(ctx.guild.id, {'always_allow_menu': not guild['always_allow_menu']}) + guild = await self.db.get_guild(ctx.guild.id, projection={'always_allow_menu': 1}) + await self.db.update(ctx.guild.id, {'always_allow_menu': not guild['always_allow_menu']}) await ctx.respond(f"Меню проигрывателя теперь {'можно' if not guild['always_allow_menu'] else 'нельзя'} создавать в каналах с несколькими людьми.", delete_after=15, ephemeral=True) @settings.command(name="vote", description="Настроить голосование.") @@ -73,10 +75,10 @@ class Settings(Cog): await ctx.respond("❌ У вас нет прав для выполнения этой команды.", delete_after=15, ephemeral=True) return - guild = self.db.get_guild(ctx.guild.id) + guild = await self.db.get_guild(ctx.guild.id, projection={'vote_next_track': 1, 'vote_add_track': 1, 'vote_add_album': 1, 'vote_add_artist': 1, 'vote_add_playlist': 1}) if vote_type == '-Всё': - self.db.update(ctx.guild.id, { + await self.db.update(ctx.guild.id, { 'vote_next_track': False, 'vote_add_track': False, 'vote_add_album': False, @@ -86,7 +88,7 @@ class Settings(Cog): ) response_message = "Голосование выключено." elif vote_type == '+Всё': - self.db.update(ctx.guild.id, { + await self.db.update(ctx.guild.id, { 'vote_next_track': True, 'vote_add_track': True, 'vote_add_album': True, @@ -96,19 +98,19 @@ class Settings(Cog): ) response_message = "Голосование включено." elif vote_type == 'Переключение': - self.db.update(ctx.guild.id, {'vote_next_track': not guild['vote_next_track']}) + await self.db.update(ctx.guild.id, {'vote_next_track': not guild['vote_next_track']}) response_message = "Голосование за переключение трека " + ("выключено." if guild['vote_next_track'] else "включено.") elif vote_type == 'Трек': - self.db.update(ctx.guild.id, {'vote_add_track': not guild['vote_add_track']}) + await self.db.update(ctx.guild.id, {'vote_add_track': not guild['vote_add_track']}) response_message = "Голосование за добавление трека " + ("выключено." if guild['vote_add_track'] else "включено.") elif vote_type == 'Альбом': - self.db.update(ctx.guild.id, {'vote_add_album': not guild['vote_add_album']}) + await self.db.update(ctx.guild.id, {'vote_add_album': not guild['vote_add_album']}) response_message = "Голосование за добавление альбома " + ("выключено." if guild['vote_add_album'] else "включено.") elif vote_type == 'Артист': - self.db.update(ctx.guild.id, {'vote_add_artist': not guild['vote_add_artist']}) + await self.db.update(ctx.guild.id, {'vote_add_artist': not guild['vote_add_artist']}) response_message = "Голосование за добавление артиста " + ("выключено." if guild['vote_add_artist'] else "включено.") elif vote_type == 'Плейлист': - self.db.update(ctx.guild.id, {'vote_add_playlist': not guild['vote_add_playlist']}) + await self.db.update(ctx.guild.id, {'vote_add_playlist': not guild['vote_add_playlist']}) response_message = "Голосование за добавление плейлиста " + ("выключено." if guild['vote_add_playlist'] else "включено.") await ctx.respond(response_message, delete_after=15, ephemeral=True) diff --git a/MusicBot/cogs/utils/embeds.py b/MusicBot/cogs/utils/embeds.py index bc791e6..96d2aeb 100644 --- a/MusicBot/cogs/utils/embeds.py +++ b/MusicBot/cogs/utils/embeds.py @@ -11,7 +11,7 @@ from yandex_music import Track, Album, Artist, Playlist, Label from discord import Embed async def generate_item_embed(item: Track | Album | Artist | Playlist | list[Track], vibing: bool = False) -> Embed: - """Generate item embed. list[Track] is used for likes. + """Generate item embed. list[Track] is used for likes. If vibing is True, add vibing image. Args: item (yandex_music.Track | yandex_music.Album | yandex_music.Artist | yandex_music.Playlist): Item to be processed. @@ -19,7 +19,7 @@ async def generate_item_embed(item: Track | Album | Artist | Playlist | list[Tra Returns: discord.Embed: Item embed. """ - logging.debug(f"Generating embed for type: '{type(item).__name__}'") + logging.debug(f"[EMBEDS] Generating embed for type: '{type(item).__name__}'") if isinstance(item, Track): embed = await _generate_track_embed(item) diff --git a/MusicBot/cogs/utils/voice_extension.py b/MusicBot/cogs/utils/voice_extension.py index 354dcbc..566e897 100644 --- a/MusicBot/cogs/utils/voice_extension.py +++ b/MusicBot/cogs/utils/voice_extension.py @@ -27,6 +27,11 @@ class VoiceExtension: self.users_db = BaseUsersDatabase() async def send_menu_message(self, ctx: ApplicationContext | Interaction) -> None: + """Send menu message to the channel. Delete old menu message if exists. + + Args: + ctx (ApplicationContext | Interaction): Context. + """ from MusicBot.ui import MenuView logging.info("[VC_EXT] Sending menu message") @@ -34,20 +39,23 @@ class VoiceExtension: logging.warning("[VC_EXT] Guild id not found in context inside 'create_menu'") return - guild = self.db.get_guild(ctx.guild_id) - embed = None + guild = await self.db.get_guild(ctx.guild_id, projection={'current_track': 1, 'current_menu': 1, 'vibing': 1}) if guild['current_track']: - track = cast(Track, Track.de_json( + track = cast(Track, await asyncio.to_thread( + Track.de_json, guild['current_track'], - client=YMClient() # type: ignore # Async client can be used here. + YMClient() # type: ignore # Async client can be used here. )) embed = await generate_item_embed(track, guild['vibing']) + vc = await self.get_voice_client(ctx) if vc and vc.is_paused(): embed.set_footer(text='Приостановлено') else: embed.remove_footer() + else: + embed = None if guild['current_menu']: logging.info(f"[VC_EXT] Deleting old menu message {guild['current_menu']} in guild {ctx.guild_id}") @@ -61,12 +69,12 @@ class VoiceExtension: interaction = cast(discord.Interaction, await ctx.respond(view=menu_views[ctx.guild_id], embed=embed)) response = await interaction.original_response() - self.db.update(ctx.guild_id, {'current_menu': response.id}) + await self.db.update(ctx.guild_id, {'current_menu': response.id}) logging.info(f"[VC_EXT] New menu message {response.id} created in guild {ctx.guild_id}") async def get_menu_message(self, ctx: ApplicationContext | Interaction | RawReactionActionEvent, menu_mid: int) -> discord.Message | None: - """Fetch the menu message by its id. Return the message if found, None if not. + """Fetch the menu message by its id. Return the message if found. Reset `current_menu` field in the database if not found. Args: @@ -95,24 +103,26 @@ class VoiceExtension: raise ValueError(f"Invalid context type: '{type(ctx).__name__}'.") except discord.DiscordException as e: logging.debug(f"[VC_EXT] Failed to get menu message: {e}") - self.db.update(ctx.guild_id, {'current_menu': None}) + await self.db.update(ctx.guild_id, {'current_menu': None}) return None if menu: logging.debug("[VC_EXT] Menu message found") else: logging.debug("[VC_EXT] Menu message not found. Resetting current_menu field.") - self.db.update(ctx.guild_id, {'current_menu': None}) + await self.db.update(ctx.guild_id, {'current_menu': None}) return menu async def update_menu_embed( self, ctx: ApplicationContext | Interaction | RawReactionActionEvent, - menu_mid: int, + menu_mid: int | None = None, + *, + menu_message: discord.Message | None = None, button_callback: bool = False ) -> bool: - """Update current menu message by its id. Return True if updated, False if not. + """Update current menu message by its id. Return True if updated, False otherwise. Args: ctx (ApplicationContext | Interaction): Context. @@ -137,24 +147,31 @@ class VoiceExtension: if not gid or not uid: logging.warning("[VC_EXT] Guild ID or User ID not found in context inside 'update_menu_embed'") return False + + if not menu_message: + if not menu_mid: + logging.debug("[VC_EXT] No menu message or menu message id provided") + return False + menu = await self.get_menu_message(ctx, menu_mid) + else: + menu = menu_message - menu = await self.get_menu_message(ctx, menu_mid) if not menu: return False - token = self.users_db.get_ym_token(uid) + token = await self.users_db.get_ym_token(uid) if not token: logging.debug(f"[VC_EXT] No token found for user {uid}") return False - guild = self.db.get_guild(gid) - current_track = guild['current_track'] - if not current_track: + guild = await self.db.get_guild(gid, projection={'vibing': 1, 'current_track': 1}) + + if not guild['current_track']: logging.debug("[VC_EXT] No current track found") return False track = cast(Track, Track.de_json( - current_track, + guild['current_track'], client=YMClient(token) # type: ignore # Async client can be used here. )) @@ -164,17 +181,21 @@ class VoiceExtension: if gid in menu_views: menu_views[gid].stop() menu_views[gid] = await MenuView(ctx).init() + if isinstance(ctx, Interaction) and button_callback: # If interaction from menu buttons await ctx.edit(embed=embed, view=menu_views[gid]) else: # If interaction from other buttons or commands. They should have their own response. await menu.edit(embed=embed, view=menu_views[gid]) + except discord.NotFound: logging.warning("[VC_EXT] Menu message not found") + if gid in menu_views: menu_views[gid].stop() del menu_views[gid] + return False logging.debug("[VC_EXT] Menu embed updated") @@ -194,7 +215,8 @@ class VoiceExtension: Args: ctx (ApplicationContext | Interaction): Context. type (Literal['track', 'album', 'artist', 'playlist', 'user']): Type of the item. - id (str | int): ID of the item. + id (str | int): ID of the YM item. + update_settings (bool, optional): Update vibe settings usind data from database. Defaults to False. button_callback (bool, optional): If the function is called from button callback. Defaults to False. Returns: @@ -208,18 +230,17 @@ class VoiceExtension: logging.warning("[VC_EXT] Guild ID or User ID not found in context inside 'vibe_update'") return None - user = self.users_db.get_user(uid) + user = await self.users_db.get_user(uid, projection={'ym_token': 1, 'vibe_settings': 1}) if not user['ym_token']: logging.info(f"[VC_EXT] User {uid} has no YM token") await ctx.respond("❌ Укажите токен через /account login.", ephemeral=True) - return + return None client = await self.init_ym_client(ctx, user['ym_token']) if not client: - return - - self.users_db.update(uid, {'vibe_type': type, 'vibe_id': id}) - guild = self.db.get_guild(gid) + return None + + guild = await self.db.get_guild(gid, projection={'vibing': 1, 'current_track': 1}) if not guild['vibing']: feedback = await client.rotor_station_feedback_radio_started( @@ -228,8 +249,10 @@ class VoiceExtension: timestamp=time() ) logging.debug(f"[VIBE] Radio started feedback: {feedback}") + if not feedback: + return None + tracks = await client.rotor_station_tracks(f"{type}:{id}") - self.db.update(gid, {'vibing': True}) if update_settings: settings = user['vibe_settings'] @@ -259,17 +282,23 @@ class VoiceExtension: if not tracks: logging.warning("[VIBE] Failed to get next vibe tracks") await ctx.respond("❌ Что-то пошло не так. Повторите попытку позже.", ephemeral=True) - return + return None logging.debug(f"[VIBE] Got next vibe tracks: {[track.track.title for track in tracks.sequence if track.track]}") - self.users_db.update(uid, {'vibe_batch_id': tracks.batch_id}) next_tracks = [cast(Track, track.track) for track in tracks.sequence] - self.db.update(gid, { - 'next_tracks': [track.to_dict() for track in next_tracks[1:]], - 'current_viber_id': uid + await self.users_db.update(uid, { + 'vibe_type': type, + 'vibe_id': id, + 'vibe_batch_id': tracks.batch_id }) + await self.db.update(gid, { + 'next_tracks': [track.to_dict() for track in next_tracks[1:]], + 'current_viber_id': uid, + 'vibing': True + }) + await self.stop_playing(ctx) return await self.play_track(ctx, next_tracks[0], button_callback=button_callback) @@ -286,7 +315,7 @@ class VoiceExtension: logging.warning("[VC_EXT] User or guild not found in context inside 'voice_check'") return False - token = self.users_db.get_ym_token(ctx.user.id) + token = await self.users_db.get_ym_token(ctx.user.id) if not token: logging.debug(f"[VC_EXT] No token found for user {ctx.user.id}") await ctx.respond("❌ Укажите токен через /account login.", delete_after=15, ephemeral=True) @@ -310,7 +339,7 @@ class VoiceExtension: return False if check_vibe_privilage: - guild = self.db.get_guild(ctx.guild.id) + guild = await self.db.get_guild(ctx.guild.id, projection={'current_viber_id': 1, 'vibing': 1}) member = cast(discord.Member, ctx.user) if guild['vibing'] and ctx.user.id != guild['current_viber_id'] and not member.guild_permissions.manage_channels: logging.debug("[VIBE] Context user is not the current viber") @@ -363,7 +392,7 @@ class VoiceExtension: retry: bool = False ) -> str | None: """Download ``track`` by its id and play it in the voice channel. Return track title on success. - If sound is already playing, add track id to the queue. There's no response to the context. + Sends feedback for vibe track playing. There's no response to the context. Args: ctx (ApplicationContext | Interaction): Context @@ -376,18 +405,15 @@ class VoiceExtension: Returns: str | None: Song title or None. """ - from MusicBot.ui import MenuView - gid = ctx.guild_id if isinstance(ctx, discord.RawReactionActionEvent) else ctx.guild.id if ctx.guild else None uid = ctx.user_id if isinstance(ctx, discord.RawReactionActionEvent) else ctx.user.id if ctx.user else None if not gid or not uid: logging.warning("Guild ID or User ID not found in context") return None + vc = await self.get_voice_client(ctx) if not vc else vc if not vc: - vc = await self.get_voice_client(ctx) - if not vc: - return None + return None if isinstance(ctx, Interaction): loop = ctx.client.loop @@ -400,13 +426,13 @@ class VoiceExtension: else: raise ValueError(f"Invalid context type: '{type(ctx).__name__}'.") - self.db.update(gid, {'current_track': track.to_dict()}) - guild = self.db.get_guild(gid) + await self.db.set_current_track(gid, track) + guild = await self.db.get_guild(gid, projection={'current_menu': 1, 'vibing': 1}) try: await asyncio.gather( - track.download_async(f'music/{gid}.mp3'), - self._update_menu(ctx, guild, track, menu_message, button_callback) + self._download_track(gid, track), + self.update_menu_embed(ctx, guild['current_menu'], menu_message=menu_message, button_callback=button_callback) ) except yandex_music.exceptions.TimedOutError: logging.warning(f"[VC_EXT] Timed out while downloading track '{track.title}'") @@ -419,17 +445,21 @@ class VoiceExtension: async with aiofiles.open(f'music/{gid}.mp3', "rb") as f: track_bytes = io.BytesIO(await f.read()) - song = discord.FFmpegPCMAudio(track_bytes, pipe=True, options='-vn -filter:a "volume=0.15"') - if not guild['current_menu']: - await asyncio.sleep(1) + song = discord.FFmpegPCMAudio(track_bytes, pipe=True, options='-vn -b:a 64k -filter:a "volume=0.15"') + + # Giving FFMPEG enough time to process the audio file + if not guild['vibing']: + await asyncio.sleep(0.75) + else: + await asyncio.sleep(0.25) vc.play(song, after=lambda exc: asyncio.run_coroutine_threadsafe(self.next_track(ctx, after=True), loop)) logging.info(f"[VC_EXT] Playing track '{track.title}'") - self.db.update(gid, {'is_stopped': False}) - + await self.db.update(gid, {'is_stopped': False}) + if guild['vibing']: - user = self.users_db.get_user(uid) + user = await self.users_db.get_user(uid) feedback = await cast(YMClient, track.client).rotor_station_feedback_track_started( f"{user['vibe_type']}:{user['vibe_id']}", track.id, @@ -454,16 +484,16 @@ class VoiceExtension: logging.warning("[VC_EXT] Guild ID not found in context") return - guild = self.db.get_guild(gid) + guild = await self.db.get_guild(gid, projection={'current_menu': 1, 'current_track': 1, 'vibing': 1}) if gid in menu_views: menu_views[gid].stop() del menu_views[gid] - if not vc: - vc = await self.get_voice_client(ctx) + + vc = await self.get_voice_client(ctx) if not vc else vc if vc: logging.debug("[VC_EXT] Stopping playback") - self.db.update(gid, {'current_track': None, 'is_stopped': True}) + await self.db.update(gid, {'current_track': None, 'is_stopped': True}) vc.stop() if full: @@ -472,13 +502,13 @@ class VoiceExtension: if menu: await menu.delete() - self.db.update(gid, { + await self.db.update(gid, { 'current_menu': None, 'repeat': False, 'shuffle': False, 'previous_tracks': [], 'next_tracks': [], 'vibing': False }) logging.info(f"[VOICE] Playback stopped in guild {gid}") if guild['vibing']: - user = self.users_db.get_user(uid) + user = await self.users_db.get_user(uid) token = user['ym_token'] if not token: logging.info(f"[VOICE] User {uid} has no YM token") @@ -539,8 +569,8 @@ class VoiceExtension: logging.warning("[VC_EXT] Guild ID or User ID not found in context inside 'next_track'") return None - guild = self.db.get_guild(gid) - user = self.users_db.get_user(uid) + guild = await self.db.get_guild(gid, projection={'shuffle': 1, 'repeat': 1, 'is_stopped': 1, 'current_menu': 1, 'vibing': 1, 'current_track': 1}) + user = await self.users_db.get_user(uid) if not user['ym_token']: logging.debug(f"[VC_EXT] No token found for user {uid}") return None @@ -602,14 +632,14 @@ class VoiceExtension: next_track = guild['current_track'] elif guild['shuffle']: logging.debug("[VC_EXT] Shuffling tracks") - next_track = self.db.get_random_track(gid) + next_track = await self.db.pop_random_track(gid, 'next') else: logging.debug("[VC_EXT] Getting next track") - next_track = self.db.get_track(gid, 'next') + next_track = await self.db.get_track(gid, 'next') if guild['current_track'] and guild['current_menu'] and not guild['repeat']: logging.debug("[VC_EXT] Adding current track to history") - self.db.modify_track(gid, guild['current_track'], 'previous', 'insert') + await self.db.modify_track(gid, guild['current_track'], 'previous', 'insert') if next_track: ym_track = Track.de_json( @@ -643,8 +673,9 @@ class VoiceExtension: button_callback=button_callback ) - logging.info("No next track found") - self.db.update(gid, {'is_stopped': True, 'current_track': None}) + logging.info("[VIBE] No next track found") + if after: + await self.db.update(gid, {'is_stopped': True, 'current_track': None}) return None async def prev_track(self, ctx: ApplicationContext | Interaction, button_callback: bool = False) -> str | None: @@ -663,9 +694,10 @@ class VoiceExtension: return None gid = ctx.guild.id - token = self.users_db.get_ym_token(ctx.user.id) - current_track = self.db.get_track(gid, 'current') - prev_track = self.db.get_track(gid, 'previous') + token = await self.users_db.get_ym_token(ctx.user.id) + current_track = await self.db.get_track(gid, 'current') + prev_track = await self.db.get_track(gid, 'previous') + print(prev_track) if not token: logging.debug(f"[VC_EXT] No token found for user {ctx.user.id}") @@ -676,7 +708,7 @@ class VoiceExtension: track: dict[str, Any] | None = prev_track elif current_track: logging.debug("[VC_EXT] No previous track found. Repeating current track") - track = self.db.get_track(gid, 'current') + track = current_track else: logging.debug("[VC_EXT] No previous or current track found") track = None @@ -711,8 +743,8 @@ class VoiceExtension: logging.warning("Guild ID or User ID not found in context inside 'play_track'") return None - current_track = self.db.get_track(gid, 'current') - token = self.users_db.get_ym_token(uid) + current_track = await self.db.get_track(gid, 'current') + token = await self.users_db.get_ym_token(uid) if not token: logging.debug(f"[VC_EXT] No token found for user {uid}") return None @@ -741,8 +773,8 @@ class VoiceExtension: logging.warning("[VC_EXT] Guild or User not found in context inside 'like_track'") return None - current_track = self.db.get_track(ctx.guild.id, 'current') - token = self.users_db.get_ym_token(ctx.user.id) + current_track = await self.db.get_track(ctx.guild.id, 'current') + token = await self.users_db.get_ym_token(ctx.user.id) if not current_track or not token: logging.debug("[VC_EXT] Current track or token not found in 'like_track'") return None @@ -782,7 +814,7 @@ class VoiceExtension: logging.warning("[VC_EXT] Guild or User not found in context inside 'dislike_track'") return False - current_track = self.db.get_track(ctx.guild.id, 'current') + current_track = await self.db.get_track(ctx.guild.id, 'current') if not current_track: logging.debug("[VC_EXT] Current track not found in 'dislike_track'") return False @@ -797,42 +829,6 @@ class VoiceExtension: ) return res - async def _retry_update_menu_embed( - self, - ctx: ApplicationContext | Interaction, - menu_mid: int, - button_callback: bool - ) -> None: - update = await self.update_menu_embed(ctx, menu_mid, button_callback) - for _ in range(10): - if update: - break - await asyncio.sleep(0.25) - update = await self.update_menu_embed(ctx, menu_mid, button_callback) - - async def _update_menu( - self, - ctx: ApplicationContext | Interaction | RawReactionActionEvent, - guild: ExplicitGuild, - track: Track, - menu_message: discord.Message | None, - button_callback: bool - ) -> None: - from MusicBot.ui import MenuView - gid = cast(int, ctx.guild_id) - - if guild['current_menu'] and not isinstance(ctx, RawReactionActionEvent): - if menu_message: - try: - if gid in menu_views: - menu_views[gid].stop() - menu_views[gid] = await MenuView(ctx).init() - await menu_message.edit(embed=await generate_item_embed(track, guild['vibing']), view=menu_views[gid]) - except discord.errors.NotFound: - logging.warning("[VC_EXT] Menu message not found. Using 'update_menu_embed' instead.") - await self._retry_update_menu_embed(ctx, guild['current_menu'], button_callback) - else: - await self._retry_update_menu_embed(ctx, guild['current_menu'], button_callback) async def init_ym_client(self, ctx: ApplicationContext | Interaction | RawReactionActionEvent, token: str | None = None) -> YMClient | None: """Initialize Yandex Music client. Return client on success. Return None if no token found and respond to the context. @@ -847,7 +843,7 @@ class VoiceExtension: if not token: uid = ctx.user_id if isinstance(ctx, discord.RawReactionActionEvent) else ctx.user.id if ctx.user else None - token = self.users_db.get_ym_token(uid) if uid else None + token = await self.users_db.get_ym_token(uid) if uid else None if not token: logging.debug("No token found in 'init_ym_client'") @@ -855,6 +851,12 @@ class VoiceExtension: await ctx.respond("❌ Укажите токен через /account login.", delete_after=15, ephemeral=True) return None + if not hasattr(self, '_ym_clients'): + self._ym_clients = {} + + if token in self._ym_clients: + return self._ym_clients[token] + try: client = await YMClient(token).init() except yandex_music.exceptions.UnauthorizedError: @@ -862,4 +864,26 @@ class VoiceExtension: if not isinstance(ctx, discord.RawReactionActionEvent): await ctx.respond("❌ Недействительный токен. Если это не так, попробуйте ещё раз.", delete_after=15, ephemeral=True) return None + + self._ym_clients[token] = client return client + + async def _retry_update_menu_embed( + self, + ctx: ApplicationContext | Interaction, + menu_mid: int, + button_callback: bool + ) -> None: + update = await self.update_menu_embed(ctx, menu_mid, button_callback=button_callback) + for _ in range(10): + if update: + break + await asyncio.sleep(0.25) + update = await self.update_menu_embed(ctx, menu_mid, button_callback=button_callback) + + async def _download_track(self, gid: int, track: Track) -> None: + try: + await track.download_async(f'music/{gid}.mp3') + except yandex_music.exceptions.TimedOutError: + logging.warning(f"[VC_EXT] Timeout downloading {track.title}") + raise diff --git a/MusicBot/cogs/voice.py b/MusicBot/cogs/voice.py index 72786c6..e1a7b7c 100644 --- a/MusicBot/cogs/voice.py +++ b/MusicBot/cogs/voice.py @@ -25,7 +25,7 @@ class Voice(Cog, VoiceExtension): logging.info(f"[VOICE] Voice state update for member {member.id} in guild {member.guild.id}") gid = member.guild.id - guild = self.db.get_guild(gid) + guild = await self.db.get_guild(gid, projection={'current_menu': 1, 'always_allow_menu': 1}) discord_guild = await self.typed_bot.fetch_guild(gid) current_menu = guild['current_menu'] @@ -46,13 +46,13 @@ class Voice(Cog, VoiceExtension): menu_views[member.guild.id].stop() del menu_views[member.guild.id] - self.db.update(gid, {'previous_tracks': [], 'next_tracks': [], 'vibing': False, 'current_menu': None, 'repeat': False, 'shuffle': False}) + await self.db.update(gid, {'previous_tracks': [], 'next_tracks': [], 'vibing': False, 'current_menu': None, 'repeat': False, 'shuffle': False}) vc.stop() elif len(channel.members) > 2 and not guild['always_allow_menu']: if current_menu: logging.info(f"[VOICE] Disabling current menu for guild {gid} due to multiple members") - self.db.update(gid, {'current_menu': None, 'repeat': False, 'shuffle': False, 'vibing': False}) + await self.db.update(gid, {'current_menu': None, 'repeat': False, 'shuffle': False, 'vibing': False}) try: message = await channel.fetch_message(current_menu) await message.delete() @@ -82,7 +82,7 @@ class Voice(Cog, VoiceExtension): if not message or message.author.id != bot_id: return - if not self.users_db.get_ym_token(payload.user_id): + if not await self.users_db.get_ym_token(payload.user_id): await message.remove_reaction(payload.emoji, payload.member) await channel.send("Для участия в голосовании необходимо авторизоваться через /account login.", delete_after=15) return @@ -90,7 +90,7 @@ class Voice(Cog, VoiceExtension): guild_id = payload.guild_id if not guild_id: return - guild = self.db.get_guild(guild_id) + guild = await self.db.get_guild(guild_id, projection={'votes': 1, 'current_track': 1}) votes = guild['votes'] if payload.message_id not in votes: @@ -113,7 +113,7 @@ class Voice(Cog, VoiceExtension): if vote_data['action'] == 'next': logging.info(f"[VOICE] Skipping track for message {payload.message_id}") - self.db.update(guild_id, {'is_stopped': False}) + await self.db.update(guild_id, {'is_stopped': False}) title = await self.next_track(payload) await message.clear_reactions() await message.edit(content=f"Сейчас играет: **{title}**!", delete_after=15) @@ -128,8 +128,8 @@ class Voice(Cog, VoiceExtension): logging.info(f"[VOICE] Recieved empty vote context for message {payload.message_id}") return - self.db.update(guild_id, {'is_stopped': False}) - self.db.modify_track(guild_id, track, 'next', 'append') + await self.db.update(guild_id, {'is_stopped': False}) + await self.db.modify_track(guild_id, track, 'next', 'append') if guild['current_track']: await message.edit(content=f"Трек был добавлен в очередь!", delete_after=15) @@ -149,8 +149,8 @@ class Voice(Cog, VoiceExtension): logging.info(f"[VOICE] Recieved empty vote context for message {payload.message_id}") return - self.db.update(guild_id, {'is_stopped': False}) - self.db.modify_track(guild_id, tracks, 'next', 'extend') + await self.db.update(guild_id, {'is_stopped': False}) + await self.db.modify_track(guild_id, tracks, 'next', 'extend') if guild['current_track']: await message.edit(content=f"Контент был добавлен в очередь!", delete_after=15) @@ -166,7 +166,7 @@ class Voice(Cog, VoiceExtension): await message.edit(content='Запрос был отклонён.', delete_after=15) del votes[str(payload.message_id)] - self.db.update(guild_id, {'votes': votes}) + await self.db.update(guild_id, {'votes': votes}) @Cog.listener() async def on_raw_reaction_remove(self, payload: discord.RawReactionActionEvent) -> None: @@ -177,7 +177,7 @@ class Voice(Cog, VoiceExtension): guild_id = payload.guild_id if not guild_id: return - guild = self.db.get_guild(guild_id) + guild = await self.db.get_guild(guild_id, projection={'votes': 1}) votes = guild['votes'] channel = cast(discord.VoiceChannel, self.typed_bot.get_channel(payload.channel_id)) @@ -196,13 +196,15 @@ class Voice(Cog, VoiceExtension): logging.info(f"[VOICE] User {payload.user_id} removed negative vote for message {payload.message_id}") del vote_data['negative_votes'][payload.user_id] - self.db.update(guild_id, {'votes': votes}) + await self.db.update(guild_id, {'votes': votes}) @voice.command(name="menu", description="Создать меню проигрывателя. Доступно только если вы единственный в голосовом канале.") async def menu(self, ctx: discord.ApplicationContext) -> None: logging.info(f"[VOICE] Menu command invoked by user {ctx.author.id} in guild {ctx.guild.id}") + if not await self.voice_check(ctx): + return - guild = self.db.get_guild(ctx.guild.id) + guild = await self.db.get_guild(ctx.guild.id, projection={'always_allow_menu': 1}) channel = cast(discord.VoiceChannel, ctx.channel) if len(channel.members) > 2 and not guild['always_allow_menu']: @@ -257,7 +259,7 @@ class Voice(Cog, VoiceExtension): logging.info(f"[VOICE] User {ctx.author.id} does not have permissions to execute leave command in guild {ctx.guild.id}") await ctx.respond("❌ У вас нет прав для выполнения этой команды.", delete_after=15, ephemeral=True) elif await self.voice_check(ctx): - self.db.update(ctx.guild.id, {'previous_tracks': [], 'next_tracks': []}) + await self.db.update(ctx.guild.id, {'previous_tracks': [], 'next_tracks': []}) await ctx.respond("Очередь и история сброшены.", delete_after=15, ephemeral=True) logging.info(f"[VOICE] Queue and history cleared in guild {ctx.guild.id}") @@ -267,11 +269,11 @@ class Voice(Cog, VoiceExtension): if not await self.voice_check(ctx): return - self.users_db.update(ctx.user.id, {'queue_page': 0}) + await self.users_db.update(ctx.user.id, {'queue_page': 0}) - tracks = self.db.get_tracks_list(ctx.guild.id, 'next') + tracks = await self.db.get_tracks_list(ctx.guild.id, 'next') embed = generate_queue_embed(0, tracks) - await ctx.respond(embed=embed, view=QueueView(ctx), ephemeral=True) + await ctx.respond(embed=embed, view=await QueueView(ctx).init(), ephemeral=True) logging.info(f"[VOICE] Queue embed sent to user {ctx.author.id} in guild {ctx.guild.id}") @@ -290,7 +292,7 @@ class Voice(Cog, VoiceExtension): if not vc.is_paused(): vc.pause() - menu = self.db.get_current_menu(ctx.guild.id) + menu = await self.db.get_current_menu(ctx.guild.id) if menu: await self.update_menu_embed(ctx, menu) @@ -314,7 +316,7 @@ class Voice(Cog, VoiceExtension): elif await self.voice_check(ctx) and (vc := await self.get_voice_client(ctx)): if vc.is_paused(): vc.resume() - menu = self.db.get_current_menu(ctx.guild.id) + menu = await self.db.get_current_menu(ctx.guild.id) if menu: await self.update_menu_embed(ctx, menu) logging.info(f"[VOICE] Track resumed in guild {ctx.guild.id}") @@ -346,7 +348,7 @@ class Voice(Cog, VoiceExtension): return gid = ctx.guild.id - guild = self.db.get_guild(gid) + guild = await self.db.get_guild(gid, projection={'next_tracks': 1, 'vote_next_track': 1}) if not guild['next_tracks']: logging.info(f"[VOICE] No tracks in queue in guild {ctx.guild.id}") await ctx.respond("❌ Нет песенен в очереди.", delete_after=15, ephemeral=True) @@ -364,7 +366,7 @@ class Voice(Cog, VoiceExtension): await response.add_reaction('✅') await response.add_reaction('❌') - self.db.update_vote( + await self.db.update_vote( gid, response.id, { @@ -378,7 +380,7 @@ class Voice(Cog, VoiceExtension): else: logging.info(f"[VOICE] Skipping vote for user {ctx.author.id} in guild {ctx.guild.id}") - self.db.update(gid, {'is_stopped': False}) + await self.db.update(gid, {'is_stopped': False}) title = await self.next_track(ctx) await ctx.respond(f"Сейчас играет: **{title}**!", delete_after=15) @@ -412,7 +414,7 @@ class Voice(Cog, VoiceExtension): if not await self.voice_check(ctx): return - guild = self.db.get_guild(ctx.guild.id) + guild = await self.db.get_guild(ctx.guild.id, projection={'always_allow_menu': 1, 'current_track': 1}) channel = cast(discord.VoiceChannel, ctx.channel) if len(channel.members) > 2 and not guild['always_allow_menu']: @@ -425,7 +427,9 @@ class Voice(Cog, VoiceExtension): return await self.send_menu_message(ctx) - await self.update_vibe(ctx, 'track', guild['current_track']['id']) + feedback = await self.update_vibe(ctx, 'track', guild['current_track']['id']) + if not feedback: + await ctx.respond("❌ Операция не удалась. Возможно, у вес нет подписки на Яндекс Музыку.", ephemeral=True) @voice.command(name='vibe', description="Запустить Мою Волну.") async def user_vibe(self, ctx: discord.ApplicationContext) -> None: @@ -433,7 +437,7 @@ class Voice(Cog, VoiceExtension): if not await self.voice_check(ctx): return - guild = self.db.get_guild(ctx.guild.id) + guild = await self.db.get_guild(ctx.guild.id, projection={'always_allow_menu': 1}) channel = cast(discord.VoiceChannel, ctx.channel) if len(channel.members) > 2 and not guild['always_allow_menu']: @@ -442,4 +446,6 @@ class Voice(Cog, VoiceExtension): return await self.send_menu_message(ctx) - await self.update_vibe(ctx, 'user', 'onyourwave') + feedback = await self.update_vibe(ctx, 'user', 'onyourwave') + if not feedback: + await ctx.respond("❌ Операция не удалась. Возможно, у вес нет подписки на Яндекс Музыку.", ephemeral=True) diff --git a/MusicBot/database/__init__.py b/MusicBot/database/__init__.py index 9612837..e0cdaab 100644 --- a/MusicBot/database/__init__.py +++ b/MusicBot/database/__init__.py @@ -1,4 +1,4 @@ -from .base import BaseGuildsDatabase, BaseUsersDatabase +from .base import BaseGuildsDatabase, BaseUsersDatabase, guilds, users from .extensions import VoiceGuildsDatabase from .user import User, ExplicitUser @@ -12,5 +12,7 @@ __all__ = [ 'ExplicitUser', 'Guild', 'ExplicitGuild', - 'MessageVotes' + 'MessageVotes', + 'guilds', + 'users', ] \ No newline at end of file diff --git a/MusicBot/database/base.py b/MusicBot/database/base.py index 9466e28..943ac30 100644 --- a/MusicBot/database/base.py +++ b/MusicBot/database/base.py @@ -1,186 +1,112 @@ -"""This documents initialises databse and contains methods to access it.""" - -from typing import Any, cast - -from pymongo import MongoClient -from pymongo.collection import Collection +from typing import Iterable, Any, cast +from pymongo import AsyncMongoClient, ReturnDocument +from pymongo.asynchronous.collection import AsyncCollection +from pymongo.results import UpdateResult from .user import User, ExplicitUser from .guild import Guild, ExplicitGuild, MessageVotes -client: MongoClient = MongoClient("mongodb://localhost:27017/") -users: Collection[ExplicitUser] = client.YandexMusicBot.users -guilds: Collection[ExplicitGuild] = client.YandexMusicBot.guilds +client: AsyncMongoClient = AsyncMongoClient("mongodb://localhost:27017/") + +db = client.YandexMusicBot +users: AsyncCollection[ExplicitUser] = db.users +guilds: AsyncCollection[ExplicitGuild] = db.guilds class BaseUsersDatabase: + DEFAULT_USER = ExplicitUser( + _id=0, + ym_token=None, + playlists=[], + playlists_page=0, + queue_page=0, + vibe_batch_id=None, + vibe_type=None, + vibe_id=None, + vibe_settings={ + 'mood': 'all', + 'diversity': 'default', + 'lang': 'any' + } + ) - def create_record(self, uid: int) -> None: - """Create user database record. - - Args: - uid (int): User id. - """ - uid = uid - users.insert_one(ExplicitUser( - _id=uid, - ym_token=None, - playlists=[], - playlists_page=0, - queue_page=0, - vibe_batch_id=None, - vibe_type=None, - vibe_id=None, - vibe_settings={ - 'mood': 'all', - 'diversity': 'default', - 'lang': 'any' - } - )) - - def update(self, uid: int, data: User | dict[Any, Any]) -> None: - """Update user record. - - Args: - uid (int): User id. - data (User | dict[Any, Any]): Updated data. - """ - self.get_user(uid) - users.update_one({'_id': uid}, {"$set": data}) - - def get_user(self, uid: int) -> ExplicitUser: - """Get user record from database. Create new entry if not present. - - Args: - uid (int): User id. - - Returns: - User: User record. - """ - user = users.find_one({'_id': uid}) - if not user: - self.create_record(uid) - user = users.find_one({'_id': uid}) - user = cast(ExplicitUser, user) - existing_fields = user.keys() - fields: ExplicitUser = ExplicitUser( - _id=0, - ym_token=None, - playlists=[], - playlists_page=0, - queue_page=0, - vibe_batch_id=None, - vibe_type=None, - vibe_id=None, - vibe_settings={ - 'mood': 'all', - 'diversity': 'default', - 'lang': 'any' - } + async def update(self, uid: int, data: User | dict[str, Any]) -> UpdateResult: + return await users.update_one( + {'_id': uid}, + {'$set': data}, + upsert=True + ) + + async def get_user(self, uid: int, projection: User | Iterable[str] | None = None) -> ExplicitUser: + user = await users.find_one_and_update( + {'_id': uid}, + {'$setOnInsert': self.DEFAULT_USER}, + return_document=ReturnDocument.AFTER, + upsert=True, + projection=projection + ) + return cast(ExplicitUser, user) + + async def get_ym_token(self, uid: int) -> str | None: + user = await users.find_one( + {'_id': uid}, + projection={'ym_token': 1} + ) + return cast(str | None, user.get('ym_token') if user else None) + + async def add_playlist(self, uid: int, playlist_data: dict) -> UpdateResult: + return await users.update_one( + {'_id': uid}, + {'$push': {'playlists': playlist_data}} ) - for field, default_value in fields.items(): - if field not in existing_fields: - user[field] = default_value - users.update_one({'_id': uid}, {"$set": {field: default_value}}) - - return user - def get_ym_token(self, uid: int) -> str | None: - user = users.find_one({'_id': uid}) - if not user: - self.create_record(uid) - user = users.find_one({'_id': uid}) - return cast(ExplicitUser, user)['ym_token'] class BaseGuildsDatabase: - - def create_record(self, gid: int) -> None: - """Create guild database record. + DEFAULT_GUILD = ExplicitGuild( + _id=0, + next_tracks=[], + previous_tracks=[], + current_track=None, + current_menu=None, + is_stopped=True, + allow_explicit=True, + always_allow_menu=False, + vote_next_track=True, + vote_add_track=True, + vote_add_album=True, + vote_add_artist=True, + vote_add_playlist=True, + shuffle=False, + repeat=False, + votes={}, + vibing=False, + current_viber_id=None + ) - Args: - gid (int): Guild id. - """ - guilds.insert_one(ExplicitGuild( - _id=gid, - next_tracks=[], - previous_tracks=[], - current_track=None, - current_menu=None, - is_stopped=True, - allow_explicit=True, - always_allow_menu=False, - vote_next_track=True, - vote_add_track=True, - vote_add_album=True, - vote_add_artist=True, - vote_add_playlist=True, - shuffle=False, - repeat=False, - votes={}, - vibing=False, - current_viber_id=None - )) - - def update(self, gid: int, data: Guild) -> None: - """Update guild record. - - Args: - gid (int): Guild id. - data (dict[Any, Any]): Updated data. - """ - self.get_guild(gid) - guilds.update_one({'_id': gid}, {"$set": data}) - - def get_guild(self, gid: int) -> ExplicitGuild: - """Get guild record from database. Create new entry if not present. - - Args: - uid (int): User id. - - Returns: - Guild: Guild record. - """ - guild = guilds.find_one({'_id': gid}) - if not guild: - self.create_record(gid) - guild = guilds.find_one({'_id': gid}) - - guild = cast(ExplicitGuild, guild) - existing_fields = guild.keys() - fields = ExplicitGuild( - _id=0, - next_tracks=[], - previous_tracks=[], - current_track=None, - current_menu=None, - is_stopped=True, - allow_explicit=True, - always_allow_menu=False, - vote_next_track=True, - vote_add_track=True, - vote_add_album=True, - vote_add_artist=True, - vote_add_playlist=True, - shuffle=False, - repeat=False, - votes={}, - vibing=False, - current_viber_id=None + async def update(self, gid: int, data: Guild | dict[str, Any]) -> UpdateResult: + return await guilds.update_one( + {'_id': gid}, + {'$set': data}, + upsert=True + ) + + async def get_guild(self, gid: int, projection: Guild | Iterable[str] | None = None) -> ExplicitGuild: + guild = await guilds.find_one_and_update( + {'_id': gid}, + {'$setOnInsert': self.DEFAULT_GUILD}, + return_document=ReturnDocument.AFTER, + upsert=True, + projection=projection + ) + return cast(ExplicitGuild, guild) + + async def update_vote(self, gid: int, mid: int, data: MessageVotes) -> UpdateResult: + return await guilds.update_one( + {'_id': gid}, + {'$set': {f'votes.{mid}': data}} + ) + + async def clear_queue(self, gid: int) -> UpdateResult: + return await guilds.update_one( + {'_id': gid}, + {'$set': {'next_tracks': []}} ) - for field, default_value in fields.items(): - if field not in existing_fields: - guild[field] = default_value - guilds.update_one({'_id': gid}, {"$set": {field: default_value}}) - - return guild - - def update_vote(self, gid: int, mid: int, data: MessageVotes) -> None: - """Update vote for a message in a guild. - - Args: - gid (int): Guild id. - mid (int): Message id. - vote (bool): Vote value. - """ - guild = self.get_guild(gid) - guild['votes'][str(mid)] = data - guilds.update_one({'_id': gid}, {"$set": {'votes': guild['votes']}}) \ No newline at end of file diff --git a/MusicBot/database/extensions.py b/MusicBot/database/extensions.py index af6c34c..2fff31b 100644 --- a/MusicBot/database/extensions.py +++ b/MusicBot/database/extensions.py @@ -1,145 +1,222 @@ from random import randint from typing import Any, Literal from yandex_music import Track +from pymongo import UpdateOne, ReturnDocument +from pymongo.errors import DuplicateKeyError -from MusicBot.database import BaseGuildsDatabase +from MusicBot.database import BaseGuildsDatabase, guilds class VoiceGuildsDatabase(BaseGuildsDatabase): - def get_tracks_list(self, gid: int, type: Literal['next', 'previous']) -> list[dict[str, Any]]: - """Get tracks list with given type. + async def get_tracks_list(self, gid: int, list_type: Literal['next', 'previous']) -> list[dict[str, Any]]: + if list_type not in ('next', 'previous'): + raise ValueError("list_type must be either 'next' or 'previous'") + projection = {f"{list_type}_tracks": 1} + guild = await self.get_guild(gid, projection=projection) + return guild.get(f"{list_type}_tracks", []) - Args: - gid (int): Guild id. - type (Literal['current', 'next', 'previous']): Track type. - - Returns: - dict[str, Any] | None: Dictionary covertable to yandex_musci.Track or None - """ - guild = self.get_guild(gid) - if type == 'next': - tracks = guild['next_tracks'] - elif type == 'previous': - tracks = guild['previous_tracks'] + async def get_track(self, gid: int, list_type: Literal['next', 'previous', 'current']) -> dict[str, Any] | None: + if list_type not in ('next', 'previous', 'current'): + raise ValueError("list_type must be either 'next' or 'previous'") - return tracks - - def get_track(self, gid: int, type: Literal['current', 'next', 'previous']) -> dict[str, Any] | None: - """Get track with given type. Pop the track from list if `type` is 'next' or 'previous'. - - Args: - gid (int): Guild id. - type (Literal['current', 'next', 'previous']): Track type. - - Returns: - dict[str, Any] | None: Dictionary covertable to yandex_musci.Track or None - """ - guild = self.get_guild(gid) - if type == 'current': - track = guild['current_track'] - elif type == 'next': - tracks = guild['next_tracks'] - if not tracks: - return None - track = tracks.pop(0) - self.update(gid, {'next_tracks': tracks}) - elif type == 'previous': - tracks = guild['previous_tracks'] - if not tracks: - return None - track = tracks.pop(0) - current_track = guild['current_track'] - if current_track: - self.modify_track(gid, current_track, 'next', 'insert') - self.update(gid, {'previous_tracks': tracks}) + if list_type == 'current': + return (await self.get_guild(gid, projection={'current_track': 1}))['current_track'] + field = f'{list_type}_tracks' + update = {'$pop': {field: -1}} + result = await guilds.find_one_and_update( + {'_id': gid}, + update, + projection={field: 1}, + return_document=ReturnDocument.BEFORE + ) + + res = result.get(field, [])[0] if result and result.get(field) else None + + if field == 'previous_tracks' and res: + await guilds.find_one_and_update( + {'_id': gid}, + {'$push': {'next_tracks': {'$each': [res], '$position': 0}}}, + projection={'next_tracks': 1} + ) + + return res + + async def modify_track( + self, + gid: int, + track: Track | dict[str, Any] | list[dict[str, Any]] | list[Track], + list_type: Literal['next', 'previous'], + operation: Literal['insert', 'append', 'extend', 'pop_start', 'pop_end'] + ) -> dict[str, Any] | None: + field = f"{list_type}_tracks" + track_data = self._normalize_track_data(track) + + operations = { + 'insert': {'$push': {field: {'$each': track_data, '$position': 0}}}, + 'append': {'$push': {field: {'$each': track_data}}}, + 'extend': {'$push': {field: {'$each': track_data}}}, + 'pop_start': {'$pop': {field: -1}}, + 'pop_end': {'$pop': {field: 1}} + } + + update = operations[operation] + try: + await guilds.update_one( + {'_id': gid}, + update, + array_filters=None + ) + return await self._get_popped_track(gid, field, operation) + except DuplicateKeyError: + await self._handle_duplicate_error(gid, field) + return await self.modify_track(gid, track, list_type, operation) + + def _normalize_track_data(self, track: Track | dict | list) -> list[dict]: + if not isinstance(track, list): + track = [track] + + return [ + t.to_dict() if isinstance(t, Track) else t + for t in track + ] + + async def pop_random_track(self, gid: int, field: Literal['next', 'previous']) -> dict[str, Any] | None: + tracks = await self.get_tracks_list(gid, field) + track = tracks.pop(randint(0, len(tracks) - 1)) if tracks else None + await self.update(gid, {f"{field}_tracks": tracks}) return track - def modify_track( - self, gid: int, - track: Track | dict[str, Any] | list[dict[str, Any]] | list[Track], - type: Literal['next', 'previous'], - operation: Literal['insert', 'append', 'extend', 'pop_start', 'pop_end', 'pop_random'] - ) -> dict[str, Any] | None: - """Perform operation of given type on tracks list of given type. + async def get_current_menu(self, gid: int) -> int | None: + guild = await self.get_guild(gid, projection={'current_menu': 1}) + return guild['current_menu'] - Args: - gid (int): Guild id. - track (Track | dict[str, Any]): yandex_music.Track or a dictionary convertable to it. - type (Literal['current', 'next', 'previous']): List type. - operation (Literal['insert', 'append', 'pop_start', 'pop_end']): Operation type. + async def _get_popped_track(self, gid: int, field: str, operation: str) -> dict[str, Any] | None: + if operation not in ('pop_start', 'pop_end', 'pop_random'): + return None - Returns: - dict[str, Any] | None: Dictionary convertable to yandex_music.Track or None. - """ - guild = self.get_guild(gid) - - if type not in ('next', 'previous'): - raise ValueError(f"Type must be either 'next' or 'previous', not '{type}'") - explicit_type: Literal['next_tracks', 'previous_tracks'] = type + '_tracks' # type: ignore[assignment] - tracks = guild[explicit_type] - pop_track = None - - if isinstance(track, list): - tracks_list = [] - for _track in track: - if isinstance(_track, Track): - tracks_list.append(_track.to_dict()) - else: - tracks_list.append(_track) - - if operation != 'extend': - raise ValueError('Can only use extend operation on lists.') - else: - tracks.extend(tracks_list) - self.update(gid, {explicit_type: tracks}) # type: ignore - else: - if isinstance(track, Track): - track = track.to_dict() - if operation == 'insert': - if type == 'previous' and len(tracks) > 50: - tracks.pop() - tracks.insert(0, track) - elif operation == 'append': - tracks.append(track) - elif operation == 'pop_start': - pop_track = tracks.pop(0) - elif operation == 'pop_end': - pop_track = tracks.pop(-1) - elif operation == 'pop_random': - pop_track = tracks.pop(randint(0, len(tracks))) - elif operation == 'extend': - raise ValueError('Can only use extend operation on lists.') - else: - raise ValueError(f"Unknown operation '{operation}'") + guild = await self.get_guild(gid, projection={field: 1}) + tracks = guild.get(field, []) - self.update(gid, {explicit_type: tracks}) # type: ignore - - return pop_track - - def get_random_track(self, gid: int) -> dict[str, Any] | None: - """Pop random track from the queue. - - Args: - gid (int): Guild id. - - Returns: - dict[str, Any] | None: Dictionary covertable to yandex_musci.Track or None - """ - tracks = self.get_tracks_list(gid, 'next') if not tracks: return None - track = tracks.pop(randint(0, len(tracks))) - self.update(gid, {'next_tracks': tracks}) - return track - - def get_current_menu(self, gid: int) -> int | None: - """Get current menu. + + if operation == 'pop_start': + return tracks[0] + elif operation == 'pop_end': + return tracks[-1] + elif operation == 'pop_random': + return tracks[randint(0, len(tracks) - 1)] + + return None + + async def _handle_duplicate_error(self, gid: int, field: str) -> None: + """Handle duplicate key errors by cleaning up the array.""" + guild = await self.get_guild(gid, projection={field: 1}) + tracks = guild.get(field, []) - Args: - gid (int): Guild id. + if not tracks: + return + + # Remove duplicates while preserving order + unique_tracks = [] + seen = set() + for track in tracks: + track_id = track.get('id') + if track_id not in seen: + seen.add(track_id) + unique_tracks.append(track) + + await guilds.update_one( + {'_id': gid}, + {'$set': {field: unique_tracks}} + ) + + async def set_current_track(self, gid: int, track: Track | dict[str, Any]) -> None: + """Set the current track and update the previous tracks list.""" + if isinstance(track, Track): + track = track.to_dict() + + await guilds.update_one( + {'_id': gid}, + { + '$set': {'current_track': track} + } + ) + + async def clear_tracks(self, gid: int, list_type: Literal['next', 'previous']) -> None: + """Clear the specified tracks list.""" + field = f"{list_type}_tracks" + await guilds.update_one( + {'_id': gid}, + {'$set': {field: []}} + ) + + async def shuffle_tracks(self, gid: int, list_type: Literal['next', 'previous']) -> None: + """Shuffle the specified tracks list.""" + field = f"{list_type}_tracks" + guild = await self.get_guild(gid, projection={field: 1}) + tracks = guild.get(field, []) + + if not tracks: + return + + shuffled_tracks = tracks.copy() + for i in range(len(shuffled_tracks) - 1, 0, -1): + j = randint(0, i) + shuffled_tracks[i], shuffled_tracks[j] = shuffled_tracks[j], shuffled_tracks[i] + + await guilds.update_one( + {'_id': gid}, + {'$set': {field: shuffled_tracks}} + ) + + async def move_track( + self, + gid: int, + from_list: Literal['next', 'previous'], + to_list: Literal['next', 'previous'], + track_index: int + ) -> bool: + """Move a track from one list to another.""" + from_field = f"{from_list}_tracks" + to_field = f"{to_list}_tracks" - Returns: int | None: Menu message id or None if not present. - """ - guild = self.get_guild(gid) - return guild['current_menu'] \ No newline at end of file + if from_field not in ('next_tracks', 'previous_tracks') or to_field not in ('next_tracks', 'previous_tracks'): + raise ValueError(f"Invalid list type: '{from_field}'") + + guild = await guilds.find_one( + {'_id': gid}, + projection={from_field: 1, to_field: 1}, + ) + + if not guild or not guild.get(from_field) or track_index >= len(guild[from_field]): + return False + + track = guild[from_field].pop(track_index) + updates = [ + UpdateOne( + {'_id': gid}, + {'$set': {from_field: guild[from_field]}}, + ), + UpdateOne( + {'_id': gid}, + {'$push': {to_field: {'$each': [track], '$position': 0}}}, + ) + ] + + await guilds.bulk_write(updates) + return True + + async def get_track_count(self, gid: int, list_type: Literal['next', 'previous']) -> int: + """Get the count of tracks in the specified list.""" + field = f"{list_type}_tracks" + guild = await self.get_guild(gid, projection={field: 1}) + return len(guild.get(field, [])) + + async def set_current_menu(self, gid: int, menu_id: int | None) -> None: + """Set the current menu message ID.""" + await guilds.update_one( + {'_id': gid}, + {'$set': {'current_menu': menu_id}} + ) \ No newline at end of file diff --git a/MusicBot/ui/find.py b/MusicBot/ui/find.py index 39af63d..a991c5f 100644 --- a/MusicBot/ui/find.py +++ b/MusicBot/ui/find.py @@ -16,18 +16,18 @@ class PlayButton(Button, VoiceExtension): self.item = item async def callback(self, interaction: Interaction) -> None: - logging.debug(f"Callback triggered for type: '{type(self.item).__name__}'") + logging.debug(f"[FIND] Callback triggered for type: '{type(self.item).__name__}'") if not interaction.guild: - logging.warning("No guild found in PlayButton callback") + logging.warning("[FIND] No guild found in PlayButton callback") return if not await self.voice_check(interaction): - logging.debug("Voice check failed in PlayButton callback") + logging.debug("[FIND] Voice check failed in PlayButton callback") return gid = interaction.guild.id - guild = self.db.get_guild(gid) + guild = await self.db.get_guild(gid, projection={'current_track': 1, 'current_menu': 1, 'vote_add_track': 1, 'vote_add_album': 1, 'vote_add_artist': 1, 'vote_add_playlist': 1}) channel = cast(discord.VoiceChannel, interaction.channel) member = cast(discord.Member, interaction.user) action: Literal['add_track', 'add_album', 'add_artist', 'add_playlist'] @@ -41,7 +41,7 @@ class PlayButton(Button, VoiceExtension): elif isinstance(self.item, Album): album = await self.item.with_tracks_async() if not album or not album.volumes: - logging.debug("Failed to fetch album tracks in PlayButton callback") + logging.debug("[FIND] Failed to fetch album tracks in PlayButton callback") await interaction.respond("Не удалось получить треки альбома.", ephemeral=True) return @@ -53,7 +53,7 @@ class PlayButton(Button, VoiceExtension): elif isinstance(self.item, Artist): artist_tracks = await self.item.get_tracks_async() if not artist_tracks: - logging.debug("Failed to fetch artist tracks in PlayButton callback") + logging.debug("[FIND] Failed to fetch artist tracks in PlayButton callback") await interaction.respond("Не удалось получить треки артиста.", ephemeral=True) return @@ -65,7 +65,7 @@ class PlayButton(Button, VoiceExtension): elif isinstance(self.item, Playlist): short_tracks = await self.item.fetch_tracks_async() if not short_tracks: - logging.debug("Failed to fetch playlist tracks in PlayButton callback") + logging.debug("[FIND] Failed to fetch playlist tracks in PlayButton callback") await interaction.respond("❌ Не удалось получить треки из плейлиста.", delete_after=15) return @@ -77,12 +77,12 @@ class PlayButton(Button, VoiceExtension): elif isinstance(self.item, list): tracks = self.item.copy() if not tracks: - logging.debug("Empty tracks list in PlayButton callback") + logging.debug("[FIND] Empty tracks list in PlayButton callback") await interaction.respond("❌ Не удалось получить треки.", delete_after=15) return action = 'add_playlist' - vote_message = f"{member.mention} хочет добавить плейлист **** в очередь.\n\n Голосуйте за добавление." + vote_message = f"{member.mention} хочет добавить плейлист **Мне Нравится** в очередь.\n\n Голосуйте за добавление." response_message = f"Плейлист **«Мне нравится»** был добавлен в очередь." else: @@ -97,7 +97,7 @@ class PlayButton(Button, VoiceExtension): await response.add_reaction('✅') await response.add_reaction('❌') - self.db.update_vote( + await self.db.update_vote( gid, response.id, { @@ -109,22 +109,22 @@ class PlayButton(Button, VoiceExtension): } ) else: - logging.debug(f"Skipping vote for '{action}' (from PlayButton callback)") + logging.debug(f"[FIND] Skipping vote for '{action}' (from PlayButton callback)") + + current_menu = await self.get_menu_message(interaction, guild['current_menu']) if guild['current_menu'] else None if guild['current_track'] is not None: - self.db.modify_track(gid, tracks, 'next', 'extend') + logging.debug(f"[FIND] Adding tracks to queue (from PlayButton callback)") + await self.db.modify_track(gid, tracks, 'next', 'extend') else: + logging.debug(f"[FIND] Playing track (from PlayButton callback)") track = tracks.pop(0) - self.db.modify_track(gid, tracks, 'next', 'extend') + await self.db.modify_track(gid, tracks, 'next', 'extend') await self.play_track(interaction, track) response_message = f"Сейчас играет: **{track.title}**!" - - current_menu = None - if guild['current_menu']: - current_menu = await self.get_menu_message(interaction, guild['current_menu']) if current_menu and interaction.message: - logging.debug(f"Deleting interaction message {interaction.message.id}: current player {current_menu.id} found") + logging.debug(f"[FIND] Deleting interaction message {interaction.message.id}: current player {current_menu.id} found") await interaction.message.delete() else: await interaction.respond(response_message, delete_after=15) @@ -145,7 +145,7 @@ class MyVibeButton(Button, VoiceExtension): logging.warning(f"[VIBE] Guild ID is None in button callback") return - guild = self.db.get_guild(gid) + guild = await self.db.get_guild(gid) channel = cast(discord.VoiceChannel, interaction.channel) if len(channel.members) > 2 and not guild['always_allow_menu']: @@ -167,7 +167,7 @@ class MyVibeButton(Button, VoiceExtension): class ListenView(View): def __init__(self, item: Track | Album | Artist | Playlist | list[Track], *items: Item, timeout: float | None = 360, disable_on_timeout: bool = True): super().__init__(*items, timeout=timeout, disable_on_timeout=disable_on_timeout) - logging.debug(f"Creating view for type: '{type(item).__name__}'") + logging.debug(f"[FIND] Creating view for type: '{type(item).__name__}'") if isinstance(item, Track): link_app = f"yandexmusic://album/{item.albums[0].id}/track/{item.id}" @@ -195,3 +195,9 @@ class ListenView(View): self.add_item(self.button2) self.add_item(self.button3) self.add_item(self.button4) + + async def on_timeout(self) -> None: + try: + return await super().on_timeout() + except discord.NotFound: + pass diff --git a/MusicBot/ui/menu.py b/MusicBot/ui/menu.py index c31bd68..1ef87c2 100644 --- a/MusicBot/ui/menu.py +++ b/MusicBot/ui/menu.py @@ -18,8 +18,8 @@ class ToggleRepeatButton(Button, VoiceExtension): if not await self.voice_check(interaction) or not interaction.guild: return gid = interaction.guild.id - guild = self.db.get_guild(gid) - self.db.update(gid, {'repeat': not guild['repeat']}) + guild = await self.db.get_guild(gid) + await self.db.update(gid, {'repeat': not guild['repeat']}) if gid in menu_views: menu_views[gid].stop() @@ -36,8 +36,8 @@ class ToggleShuffleButton(Button, VoiceExtension): if not await self.voice_check(interaction) or not interaction.guild: return gid = interaction.guild.id - guild = self.db.get_guild(gid) - self.db.update(gid, {'shuffle': not guild['shuffle']}) + guild = await self.db.get_guild(gid) + await self.db.update(gid, {'shuffle': not guild['shuffle']}) if gid in menu_views: menu_views[gid].stop() @@ -155,8 +155,8 @@ class LyricsButton(Button, VoiceExtension): if not await self.voice_check(interaction, check_vibe_privilage=False) or not interaction.guild_id or not interaction.user: return - ym_token = self.users_db.get_ym_token(interaction.user.id) - current_track = self.db.get_track(interaction.guild_id, 'current') + ym_token = await self.users_db.get_ym_token(interaction.user.id) + current_track = await self.db.get_track(interaction.guild_id, 'current') if not current_track or not ym_token: return @@ -199,7 +199,7 @@ class MyVibeButton(Button, VoiceExtension): logging.warning('[VIBE] No guild id in button callback') return - track = self.db.get_track(interaction.guild_id, 'current') + track = await self.db.get_track(interaction.guild_id, 'current') if track: logging.info(f"[MENU] Playing vibe for track '{track["id"]}'") await self.update_vibe( @@ -248,7 +248,7 @@ class MyVibeSelect(Select, VoiceExtension): return logging.info(f"[VIBE] Settings option '{custom_id}' updated to {data_value}") - self.users_db.update(interaction.user.id, {f'vibe_settings.{custom_id}': data_value}) + await self.users_db.update(interaction.user.id, {f'vibe_settings.{custom_id}': data_value}) view = MyVibeSettingsView(interaction) view.disable_all_items() @@ -263,11 +263,14 @@ class MyVibeSettingsView(View, VoiceExtension): View.__init__(self, *items, timeout=timeout, disable_on_timeout=disable_on_timeout) VoiceExtension.__init__(self, None) - if not interaction.user: + self.interaction = interaction + + async def init(self) -> None: + if not self.interaction.user: logging.warning('[VIBE] No user in settings view') return - settings = self.users_db.get_user(interaction.user.id)['vibe_settings'] + settings = (await self.users_db.get_user(self.interaction.user.id, projection={'vibe_settings'}))['vibe_settings'] diversity_settings = settings['diversity'] diversity = [ @@ -347,7 +350,7 @@ class AddToPlaylistSelect(Select, VoiceExtension): logging.debug(f"[MENU] Add to playlist select callback: {data}") playlist = cast(Playlist, await self.ym_client.users_playlists(kind=data[0], user_id=data[1])) - current_track = self.db.get_track(interaction.guild_id, 'current') + current_track = await self.db.get_track(interaction.guild_id, 'current') if not current_track: return @@ -407,13 +410,10 @@ class MenuView(View, VoiceExtension): def __init__(self, ctx: ApplicationContext | Interaction | RawReactionActionEvent, *items: Item, timeout: float | None = 3600, disable_on_timeout: bool = False): View.__init__(self, *items, timeout=timeout, disable_on_timeout=disable_on_timeout) VoiceExtension.__init__(self, None) - if not ctx.guild_id: - return self.ctx = ctx - self.guild = self.db.get_guild(ctx.guild_id) - self.repeat_button = ToggleRepeatButton(style=ButtonStyle.success if self.guild['repeat'] else ButtonStyle.secondary, emoji='🔂', row=0) - self.shuffle_button = ToggleShuffleButton(style=ButtonStyle.success if self.guild['shuffle'] else ButtonStyle.secondary, emoji='🔀', row=0) + self.repeat_button = ToggleRepeatButton(style=ButtonStyle.secondary, emoji='🔂', row=0) + self.shuffle_button = ToggleShuffleButton(style=ButtonStyle.secondary, emoji='🔀', row=0) self.play_pause_button = PlayPauseButton(style=ButtonStyle.primary, emoji='⏯', row=0) self.next_button = NextTrackButton(style=ButtonStyle.primary, emoji='⏭', row=0) self.prev_button = PrevTrackButton(style=ButtonStyle.primary, emoji='⏮', row=0) @@ -426,6 +426,16 @@ class MenuView(View, VoiceExtension): self.vibe_settings_button = MyVibeSettingsButton(style=ButtonStyle.success, emoji='🛠', row=1) async def init(self, *, disable: bool = False) -> Self: + if not self.ctx.guild_id: + return self + + self.guild = await self.db.get_guild(self.ctx.guild_id) + + if self.guild['repeat']: + self.repeat_button.style = ButtonStyle.success + if self.guild['shuffle']: + self.shuffle_button.style = ButtonStyle.success + current_track = self.guild['current_track'] likes = await self.get_likes(self.ctx) @@ -470,7 +480,7 @@ class MenuView(View, VoiceExtension): if self.guild['current_menu']: await self.stop_playing(self.ctx) - self.db.update(self.ctx.guild_id, {'current_menu': None, 'previous_tracks': [], 'vibing': False}) + await self.db.update(self.ctx.guild_id, {'current_menu': None, 'previous_tracks': [], 'vibing': False}) message = await self.get_menu_message(self.ctx, self.guild['current_menu']) if message: await message.delete() diff --git a/MusicBot/ui/other.py b/MusicBot/ui/other.py index 61769a6..4f08a8c 100644 --- a/MusicBot/ui/other.py +++ b/MusicBot/ui/other.py @@ -1,5 +1,5 @@ from math import ceil -from typing import Any +from typing import Self, Any from discord.ui import View, Button, Item from discord import ApplicationContext, ButtonStyle, Interaction, Embed @@ -45,11 +45,11 @@ class MPNextButton(Button, VoiceExtension): async def callback(self, interaction: Interaction) -> None: if not interaction.user: return - user = self.users_db.get_user(interaction.user.id) + user = await self.users_db.get_user(interaction.user.id) page = user['playlists_page'] + 1 - self.users_db.update(interaction.user.id, {'playlists_page': page}) + await self.users_db.update(interaction.user.id, {'playlists_page': page}) embed = generate_playlists_embed(page, user['playlists']) - await interaction.edit(embed=embed, view=MyPlaylists(interaction)) + await interaction.edit(embed=embed, view=await MyPlaylists(interaction).init()) class MPPrevButton(Button, VoiceExtension): def __init__(self, **kwargs): @@ -59,31 +59,37 @@ class MPPrevButton(Button, VoiceExtension): async def callback(self, interaction: Interaction) -> None: if not interaction.user: return - user = self.users_db.get_user(interaction.user.id) + user = await self.users_db.get_user(interaction.user.id) page = user['playlists_page'] - 1 - self.users_db.update(interaction.user.id, {'playlists_page': page}) + await self.users_db.update(interaction.user.id, {'playlists_page': page}) embed = generate_playlists_embed(page, user['playlists']) - await interaction.edit(embed=embed, view=MyPlaylists(interaction)) + await interaction.edit(embed=embed, view=await MyPlaylists(interaction).init()) class MyPlaylists(View, VoiceExtension): def __init__(self, ctx: ApplicationContext | Interaction, *items: Item, timeout: float | None = 360, disable_on_timeout: bool = True): View.__init__(self, *items, timeout=timeout, disable_on_timeout=disable_on_timeout) VoiceExtension.__init__(self, None) - if not ctx.user: - return - user = self.users_db.get_user(ctx.user.id) + + self.ctx = ctx + self.next_button = MPNextButton(style=ButtonStyle.primary, emoji='▶️') + self.prev_button = MPPrevButton(style=ButtonStyle.primary, emoji='◀️') + + async def init(self) -> Self: + if not self.ctx.user: + return self + + user = await self.users_db.get_user(self.ctx.user.id) count = 10 * user['playlists_page'] - - next_button = MPNextButton(style=ButtonStyle.primary, emoji='▶️') - prev_button = MPPrevButton(style=ButtonStyle.primary, emoji='◀️') - + if not user['playlists'][count + 10:]: - next_button.disabled = True + self.next_button.disabled = True if not user['playlists'][:count]: - prev_button.disabled = True + self.prev_button.disabled = True + + self.add_item(self.prev_button) + self.add_item(self.next_button) - self.add_item(prev_button) - self.add_item(next_button) + return self class QueueNextButton(Button, VoiceExtension): def __init__(self, **kwargs): @@ -93,12 +99,12 @@ class QueueNextButton(Button, VoiceExtension): async def callback(self, interaction: Interaction) -> None: if not interaction.user or not interaction.guild: return - user = self.users_db.get_user(interaction.user.id) + user = await self.users_db.get_user(interaction.user.id) page = user['queue_page'] + 1 - self.users_db.update(interaction.user.id, {'queue_page': page}) - tracks = self.db.get_tracks_list(interaction.guild.id, 'next') + await self.users_db.update(interaction.user.id, {'queue_page': page}) + tracks = await self.db.get_tracks_list(interaction.guild.id, 'next') embed = generate_queue_embed(page, tracks) - await interaction.edit(embed=embed, view=QueueView(interaction)) + await interaction.edit(embed=embed, view=await QueueView(interaction).init()) class QueuePrevButton(Button, VoiceExtension): def __init__(self, **kwargs): @@ -108,31 +114,38 @@ class QueuePrevButton(Button, VoiceExtension): async def callback(self, interaction: Interaction) -> None: if not interaction.user or not interaction.guild: return - user = self.users_db.get_user(interaction.user.id) + user = await self.users_db.get_user(interaction.user.id) page = user['queue_page'] - 1 - self.users_db.update(interaction.user.id, {'queue_page': page}) - tracks = self.db.get_tracks_list(interaction.guild.id, 'next') + await self.users_db.update(interaction.user.id, {'queue_page': page}) + tracks = await self.db.get_tracks_list(interaction.guild.id, 'next') embed = generate_queue_embed(page, tracks) - await interaction.edit(embed=embed, view=QueueView(interaction)) + await interaction.edit(embed=embed, view=await QueueView(interaction).init()) class QueueView(View, VoiceExtension): def __init__(self, ctx: ApplicationContext | Interaction, *items: Item, timeout: float | None = 360, disable_on_timeout: bool = True): View.__init__(self, *items, timeout=timeout, disable_on_timeout=disable_on_timeout) VoiceExtension.__init__(self, None) - if not ctx.user or not ctx.guild: - return - tracks = self.db.get_tracks_list(ctx.guild.id, 'next') - user = self.users_db.get_user(ctx.user.id) + self.ctx = ctx + self.next_button = QueueNextButton(style=ButtonStyle.primary, emoji='▶️') + self.prev_button = QueuePrevButton(style=ButtonStyle.primary, emoji='◀️') + + async def init(self) -> Self: + if not self.ctx.user or not self.ctx.guild: + return self + + tracks = await self.db.get_tracks_list(self.ctx.guild.id, 'next') + user = await self.users_db.get_user(self.ctx.user.id) + count = 15 * user['queue_page'] - - next_button = QueueNextButton(style=ButtonStyle.primary, emoji='▶️') - prev_button = QueuePrevButton(style=ButtonStyle.primary, emoji='◀️') - + if not tracks[count + 15:]: - next_button.disabled = True + self.next_button.disabled = True if not tracks[:count]: - prev_button.disabled = True - - self.add_item(prev_button) - self.add_item(next_button) \ No newline at end of file + self.prev_button.disabled = True + + self.add_item(self.prev_button) + self.add_item(self.next_button) + + return self + \ No newline at end of file