impr: Optimize the number of function calls

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)

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
+        self.update(gid, {'current_track': track})

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
+    allow_menu: bool  # Is /toggle menu command available

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