diff --git a/cogs/basic.py b/cogs/basic.py index 3017b8e..47ac547 100644 --- a/cogs/basic.py +++ b/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") diff --git a/voicelink/__init__.py b/voicelink/__init__.py index c839b37..09495d8 100644 --- a/voicelink/__init__.py +++ b/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 * diff --git a/voicelink/config.py b/voicelink/config.py index 55e3c1a..ff6de26 100644 --- a/voicelink/config.py +++ b/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", {}) diff --git a/voicelink/enums.py b/voicelink/enums.py index bc92915..084cbbd 100644 --- a/voicelink/enums.py +++ b/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 """ diff --git a/voicelink/ipc/methods.py b/voicelink/ipc/methods.py index 01c60d1..886bc41 100644 --- a/voicelink/ipc/methods.py +++ b/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", diff --git a/voicelink/objects.py b/voicelink/objects.py index 5176bc2..f9e38fb 100644 --- a/voicelink/objects.py +++ b/voicelink/objects.py @@ -21,18 +21,21 @@ 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. You can also pass in commands.Context to get a discord.py Context object in your track. @@ -90,17 +93,38 @@ class Track: self.end_time: Optional[int] = None def __eq__(self, other) -> bool: - if not isinstance(other, Track): - return False - - return other.track_id == self.track_id + """Checks equality between two tracks.""" + 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" 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: diff --git a/voicelink/player.py b/voicelink/player.py index dc3e36f..2f7efc8 100644 --- a/voicelink/player.py +++ b/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) diff --git a/voicelink/pool.py b/voicelink/pool.py index 0c51068..d415a47 100644 --- a/voicelink/pool.py +++ b/voicelink/pool.py @@ -390,26 +390,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: