diff --git a/MusicBot/database/__init__.py b/MusicBot/database/__init__.py new file mode 100644 index 0000000..9d5ff12 --- /dev/null +++ b/MusicBot/database/__init__.py @@ -0,0 +1,15 @@ +from .base import BaseGuildsDatabase, BaseUsersDatabase +from .extensions import VoiceGuildsDatabase + +from .user import User, ExplicitUser +from .guild import Guild, ExplicitGuild + +__all__ = [ + 'BaseGuildsDatabase', + 'BaseUsersDatabase', + 'VoiceGuildsDatabase', + 'User', + 'ExplicitUser', + 'Guild', + 'ExplicitGuild', +] \ No newline at end of file diff --git a/MusicBot/database/base.py b/MusicBot/database/base.py index 05fc6f9..8d7197b 100644 --- a/MusicBot/database/base.py +++ b/MusicBot/database/base.py @@ -5,85 +5,98 @@ from typing import Any, cast from pymongo import MongoClient from pymongo.collection import Collection -from yandex_music import Track - -from MusicBot.database.user import User, ExplicitUser, TrackInfo +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 - -def create_record(uid: int | str) -> None: - """Create user database record. +guilds: Collection[Guild] = client.YandexMusicBot.guilds - Args: - uid (int | str): User id. - """ - uid = str(uid) - users.insert_one(ExplicitUser( - _id=uid, - ym_token=None, - tracks_list=[], - current_track=None, - is_stopped=True - )) +class BaseUsersDatabase: -def update(uid: int | str, data: dict[Any, Any]) -> None: - """Update user record. + def create_record(self, uid: int) -> None: + """Create user database record. - Args: - uid (int | str): User id. - data (dict[Any, Any]): Updated data. - """ - get_user(uid) - users.update_one({'_id': str(uid)}, {"$set": data}) + Args: + uid (int): User id. + """ + uid = uid + users.insert_one(ExplicitUser( + _id=uid, + ym_token=None + )) -def get_user(uid: int | str) -> User: - """Get user record from database. Create new entry if not present. + def update(self, uid: int, data: User) -> None: + """Update user record. - Args: - uid (int | str): User id. + Args: + uid (int): User id. + data (dict[Any, Any]): Updated data. + """ + self.get_user(uid) + users.update_one({'_id': uid}, {"$set": data}) - Returns: - User: User record. - """ - user = users.find_one({'_id': str(uid)}) - if not user: - create_record(uid) - user = users.find_one({'_id': str(uid)}) - return cast(User, user) + def get_user(self, uid: int) -> User: + """Get user record from database. Create new entry if not present. -def get_ym_token(uid: int | str) -> str | None: - user = users.find_one({'_id': str(uid)}) - if not user: - create_record(uid) - user = cast(User, users.find_one({'_id': str(uid)})) - return user['ym_token'] + Args: + uid (int): User id. -def get_tracks_list(uid: int | str) -> list[TrackInfo]: - user = get_user(uid) - return user.get('tracks_list') + Returns: + User: User record. + """ + user = users.find_one({'_id': uid}) + if not user: + self.create_record(uid) + user = users.find_one({'_id': uid}) + return cast(User, user) -def pop_track(uid: int | str) -> TrackInfo: - tracks_list = get_tracks_list(uid) - track = tracks_list.pop(0) - update(uid, {'tracks_list': tracks_list}) - return track + 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'] -def add_track(uid: int | str, track: Track | TrackInfo) -> None: - tracks_list = get_tracks_list(uid) - if isinstance(track, Track): - track = TrackInfo( - track_id=str(track.id), - title=track.title, # type: ignore - avail=track.available, # type: ignore - artists=", ".join(track.artists_name()), - albums=", ".join([album.title for album in track.albums]), # type: ignore - duration=track.duration_ms, # type: ignore - explicit=track.explicit or bool(track.content_warning), - bg_video=track.background_video_uri - ) - tracks_list.append(track) - update(uid, {'tracks_list': tracks_list}) +class BaseGuildsDatabase: + + def create_record(self, gid: int) -> None: + """Create guild database record. -def set_current_track(uid: int | str) -> None: - update(uid, {'current_track': str(uid)}) \ No newline at end of file + Args: + gid (int): Guild id. + """ + guilds.insert_one(ExplicitGuild( + _id=gid, + tracks_list=[], + previous_tracks_list=[], + current_track=None, + is_stopped=True, + allow_explicit=True, + allow_menu=True + )) + + def update(self, gid: int, data: dict[Any, Any]) -> 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) -> Guild: + """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}) + return cast(Guild, guild) diff --git a/MusicBot/database/extensions.py b/MusicBot/database/extensions.py new file mode 100644 index 0000000..2ab5b62 --- /dev/null +++ b/MusicBot/database/extensions.py @@ -0,0 +1,50 @@ +from typing import Any +from yandex_music import Track + +from MusicBot.database.base import BaseGuildsDatabase + +class VoiceGuildsDatabase(BaseGuildsDatabase): + + 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_tracks_list(self, gid: int) -> list[dict[str, Any]]: + guild = self.get_guild(gid) + return guild.get('tracks_list') + + def pop_track(self, gid: int) -> dict[str, Any]: + tracks_list = self.get_tracks_list(gid) + track = tracks_list.pop(0) + self.update(gid, {'tracks_list': tracks_list}) + return track + + def insert_track(self, gid: int, track: Track | dict[str, Any]) -> None: + if isinstance(track, Track): + track = track.to_dict(for_request=True) + tracks_list = self.get_tracks_list(gid) + tracks_list.insert(0, track) + self.update(gid, {'tracks_list': tracks_list}) + + def add_track(self, gid: int, track: Track) -> None: + tracks_list = self.get_tracks_list(gid) + tracks_list.append(track.to_dict(for_request=True)) + self.update(gid, {'tracks_list': tracks_list}) + + def set_current_track(self, gid: int, track: Track) -> None: + self.update(gid, {'current_track': track.to_dict(for_request=True)}) + + def add_previous_track(self, gid: int, track: Track | dict[str, Any]) -> None: + tracks_list = self.get_previous_tracks_list(gid) + if isinstance(track, Track): + track = track.to_dict(for_request=True) + tracks_list.insert(0, track) + if len(tracks_list) > 50: + tracks_list.pop() + self.update(gid, {'previous_tracks_list': tracks_list}) + + def pop_previous_track(self, gid: int) -> dict[str, Any]: + tracks_list = self.get_previous_tracks_list(gid) + track = tracks_list.pop(0) + self.update(gid, {'previous_tracks_list': tracks_list}) + return track \ No newline at end of file diff --git a/MusicBot/database/guild.py b/MusicBot/database/guild.py index 48c1cbc..6b50099 100644 --- a/MusicBot/database/guild.py +++ b/MusicBot/database/guild.py @@ -1,10 +1,18 @@ -from typing import TypedDict +from typing import TypedDict, Any class Guild(TypedDict): + tracks_list: list[dict[str, Any]] + previous_tracks_list: list[dict[str, Any]] + current_track: dict[str, Any] | None + is_stopped: bool allow_explicit: bool allow_menu: bool class ExplicitGuild(TypedDict): - _id: str + _id: int + tracks_list: list[dict[str, Any]] + previous_tracks_list: list[dict[str, Any]] + current_track: dict[str, Any] | None + is_stopped: bool # Prevents the `after` callback of play_track allow_explicit: bool - allow_menu: bool \ No newline at end of file + allow_menu: bool # Is /toggle mnu command available \ No newline at end of file diff --git a/MusicBot/database/user.py b/MusicBot/database/user.py index 98feb7b..7bf8edc 100644 --- a/MusicBot/database/user.py +++ b/MusicBot/database/user.py @@ -1,25 +1,8 @@ from typing import TypedDict -class TrackInfo(TypedDict): - track_id: str - title: str - avail: bool - artists: str - albums: str - duration: int - explicit: bool - bg_video: str | None - class User(TypedDict): ym_token: str | None - tracks_list: list[TrackInfo] - current_track: int | None - is_stopped: bool class ExplicitUser(TypedDict): - _id: str + _id: int ym_token: str | None - tracks_list: list[TrackInfo] - current_track: int | None - is_stopped: bool # Prevents callback of play_track -