Refactor track recommendation logic and add TrackRecType enum

Moved recommendation query generation from Node to Track using the new TrackRecType enum, simplifying and centralizing platform-specific recommendation formatting. Updated related methods to use Track.get_recommendations, removed Node.get_recommendations, and improved enum handling for search platform matching. Also exposed TrackRecType in voicelink.__init__.

cogs/basic.py

@@ -226,7 +226,7 @@ class Basic(commands.Cog):
         if url(query):
             return await send_localized_message(ctx, "search.noLinkSupport", ephemeral=True)
         
-        search_type: voicelink.SearchType = voicelink.SearchType.match(platform) or Config().search_platform
+        search_type: voicelink.SearchType = voicelink.SearchType.from_platform(platform) or Config().search_platform
         tracks = await player.get_tracks(query=query, requester=ctx.author, search_type=search_type)
         if not tracks:
             return await send_localized_message(ctx, "player.errors.noTrackFound")

voicelink/__init__.py

@@ -27,7 +27,7 @@ __license__ = "MIT"
 __copyright__ = "Copyright 2023 - present (c) Vocard Development, ChocoMeow"
 
 from .config import Config
-from .enums import SearchType, LoopType
+from .enums import SearchType, LoopType, TrackRecType
 from .events import *
 from .exceptions import *
 from .filters import *

voicelink/config.py

@@ -84,7 +84,7 @@ class Config:
         self.invite_link: str = "https://discord.gg/wRCgB7vBQv"
         self.nodes: Dict[str, Dict[str, Union[str, int, bool]]] = settings.get("nodes", {})
         self.max_queue: int = settings.get("default_max_queue", 1000)
-        self.search_platform: SearchType = SearchType.match(settings.get("default_search_platform", "youtube")) or SearchType.YOUTUBE
+        self.search_platform: SearchType = SearchType.from_platform(settings.get("default_search_platform", "youtube")) or SearchType.YOUTUBE
         self.bot_prefix: str = settings.get("prefix", "")
         self.activity: List[Dict[str, str]] = settings.get("activity", [{"listen": "/help"}])
         self.logging: Dict[Union[str, Dict[str, Union[str, bool]]]] = settings.get("logging", {})

voicelink/enums.py

@@ -92,7 +92,7 @@ class SearchType(Enum):
         return self.value
     
     @classmethod
-    def match(cls, value: str):
+    def from_platform(cls, value: str):
         """find an enum based on a search string."""
         normalized_value = value.lower().replace("_", "").replace(" ", "")
 
@@ -106,6 +106,81 @@ class SearchType(Enum):
     def display_name(self) -> str:
         return self.name.replace("_", " ").title()
 
+class TrackRecType(Enum):
+    """Enum representing track recommendation key formats for various platforms.
+
+    Each key format is used to generate a recommendation link or identifier
+    for a given track ID on the respective platform.
+
+    - RecommendationType.SPOTIFY:
+      Generates a Spotify recommendation key in the format 'sprec:mix:track:{track_id}'.
+
+    - RecommendationType.YOUTUBE:
+      Generates a YouTube recommendation URL with a playlist context.
+      
+    - RecommendationType.YOUTUBE_MUSIC:
+      Same as YouTube, generates a YouTube Music recommendation URL.
+      
+    - RecommendationType.DEEZER:
+      Generates a Deezer recommendation key in the format 'dzrec:{track_id}'.
+      
+    - RecommendationType.YANDEX_MUSIC:
+      Generates a Yandex Music recommendation key in the format 'ymrec:{track_id}'.
+    
+    - RecommendationType.VK_MUSIC:
+      Generates a VK Music recommendation key in the format 'vkrec:{track_id}'.
+    
+    - RecommendationType.TIDAL:
+      Generates a Tidal recommendation key in the format 'tdrec:{track_id}'.
+    
+    - RecommendationType.QOBUZ:
+      Generates a Qobuz recommendation key in the format 'qbrec:{track_id}'.
+    
+    - RecommendationType.JIOSAAVN:
+      Generates a JioSaavn recommendation key in the format 'jsrec:{track_id}'
+    """
+
+    YOUTUBE = "https://www.youtube.com/watch?v={track_id}&list=RD{track_id}"
+    YOUTUBE_MUSIC = YOUTUBE
+    SPOTIFY = "sprec:mix:track:{track_id}"
+    DEEZER = "dzrec:{track_id}"
+    YANDEX_MUSIC = "ymrec:{track_id}"
+    VK_MUSIC = "vkrec:{track_id}"
+    TIDAL = "tdrec:{track_id}"
+    QOBUZ = "qbrec:{track_id}"
+    JIOSAAVN = "jsrec:{track_id}"
+
+    def __str__(self) -> str:
+        return self.name
+
+    def format(self, track_id: str) -> str:
+        """Format the recommendation key using the provided track ID.
+        
+        Args:
+            track_id (str): The ID of the track to format.
+        
+        Returns:
+            str: The formatted recommendation link.
+        """
+        return self.value.format(track_id=track_id)
+
+    @classmethod
+    def from_platform(cls, platform: str) -> 'TrackRecType':
+        """Find the enum member based on a platform name.
+        
+        Args:
+            platform (str): The name of the platform.
+        
+        Returns:
+            TrackRecType: The corresponding enum member, or None if not found.
+        """
+        normalized = platform.lower().replace("_", "").replace(" ", "")
+        for member in cls:
+            if member.name.lower().replace("_", "") == normalized:
+                return member
+            
+        return None
+      
 class RequestMethod(Enum):
     """The enum for the different request methods in Voicelink
     """

voicelink/ipc/methods.py

@@ -152,7 +152,7 @@ async def getRecommendation(bot: commands.Bot, data: Dict) -> None:
     
     track_data = Track.decode(track_id := data.get("trackId"))
     track = Track(track_id=track_id, info=track_data, requester=bot.user)
-    tracks: List[Track] = await node.get_recommendations(track, limit=60)
+    tracks: List[Track] = await track.get_recommendations(node)
 
     return {
         "op": "getRecommendation",

voicelink/objects.py

@@ -21,17 +21,20 @@ OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
 SOFTWARE.
 """
 
-import re
+from __future__ import annotations
 
-from typing import Optional
+from typing import Optional, List, TYPE_CHECKING
 from tldextract import extract
 from discord import Member
 
-from .enums import SearchType
+from .enums import SearchType, TrackRecType
 from .config import Config
 from .utils import format_ms
 from .transformer import encode, decode
 
+if TYPE_CHECKING:
+    from .pool import Node
+
     
 class Track:
     """The base track object. Returns critical track information needed for parsing by Lavalink.
@@ -90,17 +93,38 @@ class Track:
         self.end_time: Optional[int] = None
 
     def __eq__(self, other) -> bool:
-        if not isinstance(other, Track):
+        """Checks equality between two tracks."""
-            return False
+        if isinstance(other, Track):
-
             return other.track_id == self.track_id
         
+        return False
+
     def __str__(self) -> str:
+        """String representation of the track."""
         return self.title
 
     def __repr__(self) -> str:
         return f"<Voicelink.track title={self.title!r} uri=<{self.uri!r}> length={self.length}>"
 
+    async def get_recommendations(self, node: Node) -> List[Track]:
+        """Fetches recommended tracks based on the current track."""
+        if not node or not node._available:
+            return []
+
+        rec_type = TrackRecType.from_platform(self.source)
+        if not rec_type:
+            return []
+        
+        query = rec_type.format(track_id=self.identifier)
+        tracks = await node.get_tracks(query=query, requester=node.bot.user)
+        if not tracks:
+            return []
+        
+        if isinstance(tracks, Playlist):
+            tracks = tracks.tracks
+
+        return tracks
+
     @property
     def track_id(self) -> str:
         if not self._track_id:

voicelink/player.py

@@ -857,14 +857,14 @@ class Player(VoiceProtocol):
                 await self.set_pause(True)
     
     async def get_recommendations(self, *, track: Optional[Track] = None) -> bool:
-        """Get recommendations from Youtube or Spotify."""
+        """Fetches and adds recommended tracks based on the provided track or recent history."""
         if not track:
             try:
                 track = choice(self.queue.history(incTrack=True)[-5:])
             except IndexError:
                 return False
 
-        tracks = await self._node.get_recommendations(track)
+        tracks = await track.get_recommendations(self._node)
         if tracks:
             await self.add_track(tracks, duplicate=False)
             

voicelink/pool.py

@@ -391,26 +391,6 @@ class Node:
         elif load_type == "track":
             return [Track(track_id=data["encoded"], info=data["info"], requester=requester)]
     
-    async def get_recommendations(self, track: Track, limit: int = 20) -> List[Optional[Track]]:
-        query = ""
-        if track.source == "youtube":
-            query = f"https://www.youtube.com/watch?v={track.identifier}&list=RD{track.identifier}"
-
-        elif track.source == "spotify":
-            query = f"sprec:seed_tracks={track.identifier}"
-
-        if not query:
-            return []
-        
-        tracks = await self.get_tracks(query=query, requester=self.bot.user)
-        if not tracks:
-            return []
-        
-        if isinstance(tracks, Playlist):
-            tracks = tracks.tracks
-
-        return tracks[:limit] if limit else tracks
-    
     async def update_refresh_yt_access_token(self, token: YTToken) -> dict:
         if not self._available:
             raise NodeNotAvailable(f"The node '{self._identifier}' is unavailable.")