diff --git a/MusicBot/database/base.py b/MusicBot/database/base.py index 32afcb8..0cb2edd 100644 --- a/MusicBot/database/base.py +++ b/MusicBot/database/base.py @@ -1,6 +1,6 @@ """This documents initialises databse and contains methods to access it.""" -from typing import Any, cast +from typing import cast from pymongo import MongoClient from pymongo.collection import Collection @@ -9,8 +9,8 @@ from MusicBot.database.user import User, ExplicitUser from MusicBot.database.guild import Guild, ExplicitGuild client: MongoClient = MongoClient("mongodb://localhost:27017/") -users: Collection[User] = client.YandexMusicBot.users -guilds: Collection[Guild] = client.YandexMusicBot.guilds +users: Collection[ExplicitUser] = client.YandexMusicBot.users +guilds: Collection[ExplicitGuild] = client.YandexMusicBot.guilds class BaseUsersDatabase: @@ -36,7 +36,7 @@ class BaseUsersDatabase: self.get_user(uid) users.update_one({'_id': uid}, {"$set": data}) - def get_user(self, uid: int) -> User: + def get_user(self, uid: int) -> ExplicitUser: """Get user record from database. Create new entry if not present. Args: @@ -49,14 +49,14 @@ class BaseUsersDatabase: if not user: self.create_record(uid) user = users.find_one({'_id': uid}) - return cast(User, user) + return cast(ExplicitUser, user) def get_ym_token(self, uid: int) -> str | None: user = users.find_one({'_id': uid}) if not user: self.create_record(uid) - user = cast(User, users.find_one({'_id': uid})) - return user['ym_token'] + user = users.find_one({'_id': uid}) + return cast(ExplicitUser, user)['ym_token'] class BaseGuildsDatabase: @@ -68,15 +68,16 @@ class BaseGuildsDatabase: """ guilds.insert_one(ExplicitGuild( _id=gid, - next_tracks_list=[], - previous_tracks_list=[], + next_tracks=[], + previous_tracks=[], current_track=None, + current_player=None, is_stopped=True, allow_explicit=True, allow_menu=True )) - def update(self, gid: int, data: dict[Any, Any]) -> None: + def update(self, gid: int, data: Guild) -> None: """Update guild record. Args: @@ -86,7 +87,7 @@ class BaseGuildsDatabase: self.get_guild(gid) guilds.update_one({'_id': gid}, {"$set": data}) - def get_guild(self, gid: int) -> Guild: + def get_guild(self, gid: int) -> ExplicitGuild: """Get guild record from database. Create new entry if not present. Args: @@ -99,4 +100,4 @@ class BaseGuildsDatabase: if not guild: self.create_record(gid) guild = guilds.find_one({'_id': gid}) - return cast(Guild, guild) + return cast(ExplicitGuild, guild) diff --git a/MusicBot/database/extensions.py b/MusicBot/database/extensions.py index 851bf8c..efaa4ba 100644 --- a/MusicBot/database/extensions.py +++ b/MusicBot/database/extensions.py @@ -1,62 +1,121 @@ -from typing import Any +from typing import Any, Literal, cast from yandex_music import Track -from MusicBot.database.base import BaseGuildsDatabase +from MusicBot.database import BaseGuildsDatabase class VoiceGuildsDatabase(BaseGuildsDatabase): - def clear_queue(self, gid: int) -> None: - self.update(gid, {'next_tracks_list': []}) - def clear_history(self, gid: int) -> None: - self.update(gid, {'previous_tracks_list': []}) - - def get_current_track(self, gid: int) -> dict[str, Any] | None: - guild = self.get_guild(gid) - return guild.get('current_track') - - def get_previous_tracks_list(self, gid: int) -> list[dict[str, Any]]: - guild = self.get_guild(gid) - return guild.get('previous_tracks_list') - - def get_next_tracks_list(self, gid: int) -> list[dict[str, Any]]: - guild = self.get_guild(gid) - return guild.get('next_tracks_list') + """Clear previous and next tracks list. - def get_next_track(self, gid: int) -> dict[str, Any]: - tracks_list = self.get_next_tracks_list(gid) - track = tracks_list.pop(0) - self.update(gid, {'next_tracks_list': tracks_list}) + Args: + gid (int): _description_ + """ + self.update(gid, {'previous_tracks': [], 'next_tracks': []}) + + def get_tracks_list(self, gid: int, type: Literal['next', 'previous']) -> list[dict[str, Any]]: + """Get tracks list with given type. + + 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'] + + 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 + track = tracks.pop(0) + self.update(gid, {'next_tracks': tracks}) + elif type == 'previous': + tracks = guild['previous_tracks'] + if not tracks: + return + track = tracks.pop(0) + self.update(gid, {'previous_tracks': tracks}) + return track - def insert_track(self, gid: int, track: Track | dict[str, Any]) -> None: - if isinstance(track, Track): - track = track.to_dict() - tracks_list = self.get_next_tracks_list(gid) - tracks_list.insert(0, track) - self.update(gid, {'next_tracks_list': tracks_list}) - - def append_track(self, gid: int, track: Track) -> None: - tracks_list = self.get_next_tracks_list(gid) - tracks_list.append(track.to_dict()) - self.update(gid, {'next_tracks_list': tracks_list}) + 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'] + ) -> dict[str, Any] | None: + """Perform operation of given type on tracks list of given type. - def set_current_track(self, gid: int, track: Track) -> None: - self.update(gid, {'current_track': track.to_dict()}) - - def add_previous_track(self, gid: int, track: Track | dict[str, Any]) -> None: - tracks_list = self.get_previous_tracks_list(gid) + 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. + + Returns: + dict[str, Any] | None: Dictionary convertable to yandex_music.Track or None. + """ + guild = self.get_guild(gid) + explicit_type: Literal['next_tracks', 'previous_tracks'] = type + '_tracks' + 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(0) + elif operation == 'extend': + raise ValueError('Can only use extend operation on lists.') + + self.update(gid, {explicit_type: tracks}) # type: ignore + + if pop_track: + return pop_track + + def set_current_track(self, gid: int, track: Track | dict[str, Any]) -> None: if isinstance(track, Track): track = track.to_dict() - tracks_list.insert(0, track) - if len(tracks_list) > 50: - tracks_list.pop() - self.update(gid, {'previous_tracks_list': tracks_list}) - - def get_previous_track(self, gid: int) -> dict[str, Any] | None: - tracks_list = self.get_previous_tracks_list(gid) - if len(tracks_list) == 0: - return - track = tracks_list.pop(0) - self.update(gid, {'previous_tracks_list': tracks_list}) - return track \ No newline at end of file + self.update(gid, {'current_track': track}) \ No newline at end of file diff --git a/MusicBot/database/guild.py b/MusicBot/database/guild.py index 139b31b..05d2145 100644 --- a/MusicBot/database/guild.py +++ b/MusicBot/database/guild.py @@ -1,18 +1,20 @@ from typing import TypedDict, Any -class Guild(TypedDict): - next_tracks_list: list[dict[str, Any]] - previous_tracks_list: list[dict[str, Any]] +class Guild(TypedDict, total=False): + next_tracks: list[dict[str, Any]] + previous_tracks: list[dict[str, Any]] current_track: dict[str, Any] | None + current_player: int | None is_stopped: bool allow_explicit: bool allow_menu: bool class ExplicitGuild(TypedDict): _id: int - next_tracks_list: list[dict[str, Any]] - previous_tracks_list: list[dict[str, Any]] + next_tracks: list[dict[str, Any]] + previous_tracks: list[dict[str, Any]] current_track: dict[str, Any] | None + current_player: int | None is_stopped: bool # Prevents the `after` callback of play_track allow_explicit: bool - allow_menu: bool # Is /toggle mnu command available \ No newline at end of file + allow_menu: bool # Is /toggle menu command available \ No newline at end of file diff --git a/MusicBot/database/user.py b/MusicBot/database/user.py index 7bf8edc..e03befd 100644 --- a/MusicBot/database/user.py +++ b/MusicBot/database/user.py @@ -1,6 +1,6 @@ from typing import TypedDict -class User(TypedDict): +class User(TypedDict, total=False): ym_token: str | None class ExplicitUser(TypedDict):