2.10.0

cloudwithax · cloudwithax · commit 9bffdebe2516 · 2025-10-04T00:00:57.000-04:00
diff --git a/pomice/__init__.py b/pomice/__init__.py
@@ -20,7 +20,7 @@ class DiscordPyOutdated(Exception):
         "using 'pip install discord.py'",
     )
 
-__version__ = "2.9.2"
+__version__ = "2.10.0"
 __title__ = "pomice"
 __author__ = "cloudwithax"
 __license__ = "GPL-3.0"
diff --git a/pomice/applemusic/client.py b/pomice/applemusic/client.py
@@ -1,11 +1,14 @@
 from __future__ import annotations
 
+import asyncio
 import base64
 import logging
 import re
 from datetime import datetime
+from typing import AsyncGenerator
 from typing import Dict
 from typing import List
+from typing import Optional
 from typing import Union
 
 import aiohttp
@@ -17,10 +20,10 @@
 __all__ = ("Client",)
 
 AM_URL_REGEX = re.compile(
-    r"https?://music.apple.com/(?P<country>[a-zA-Z]{2})/(?P<type>album|playlist|song|artist)/(?P<name>.+)/(?P<id>[^?]+)",
+    r"https?://music\.apple\.com/(?P<country>[a-zA-Z]{2})/(?P<type>album|playlist|song|artist)/(?P<name>.+?)/(?P<id>[^/?]+?)(?:/)?(?:\?.*)?$",
 )
 AM_SINGLE_IN_ALBUM_REGEX = re.compile(
-    r"https?://music.apple.com/(?P<country>[a-zA-Z]{2})/(?P<type>album|playlist|song|artist)/(?P<name>.+)/(?P<id>.+)(\?i=)(?P<id2>.+)",
+    r"https?://music\.apple\.com/(?P<country>[a-zA-Z]{2})/(?P<type>album|playlist|song|artist)/(?P<name>.+)/(?P<id>[^/?]+)(\?i=)(?P<id2>[^&]+)(?:&.*)?$",
 )
 
 AM_SCRIPT_REGEX = re.compile(r'<script.*?src="(/assets/index-.*?)"')
@@ -35,12 +38,14 @@ class Client:
     and translating it to a valid Lavalink track. No client auth is required here.
     """
 
-    def __init__(self) -> None:
+    def __init__(self, *, playlist_concurrency: int = 6) -> None:
         self.expiry: datetime = datetime(1970, 1, 1)
         self.token: str = ""
         self.headers: Dict[str, str] = {}
         self.session: aiohttp.ClientSession = None  # type: ignore
         self._log = logging.getLogger(__name__)
+        # Concurrency knob for parallel playlist page retrieval
+        self._playlist_concurrency = max(1, playlist_concurrency)
 
     async def _set_session(self, session: aiohttp.ClientSession) -> None:
         self.session = session
@@ -167,25 +172,127 @@ async def search(self, query: str) -> Union[Album, Playlist, Song, Artist]:
                     "This playlist is empty and therefore cannot be queued.",
                 )
 
-            _next = track_data.get("next")
-            if _next:
-                next_page_url = AM_BASE_URL + _next
-
-                while next_page_url is not None:
-                    resp = await self.session.get(next_page_url, headers=self.headers)
+            # Apple Music uses cursor pagination with 'next'. We'll fetch subsequent pages
+            # concurrently by first collecting cursors in rolling waves.
+            next_cursor = track_data.get("next")
+            semaphore = asyncio.Semaphore(self._playlist_concurrency)
 
+            async def fetch_page(url: str) -> List[Song]:
+                async with semaphore:
+                    resp = await self.session.get(url, headers=self.headers)
                     if resp.status != 200:
-                        raise AppleMusicRequestException(
-                            f"Error while fetching results: {resp.status} {resp.reason}",
-                        )
+                        if self._log:
+                            self._log.warning(
+                                f"Apple Music page fetch failed {resp.status} {resp.reason} for {url}",
+                            )
+                        return []
+                    pj: dict = await resp.json(loads=json.loads)
+                    songs = [Song(track) for track in pj.get("data", [])]
+                    # Return songs; we will look for pj.get('next') in streaming iterator variant
+                    return songs, pj.get("next")  # type: ignore
+
+            # We'll implement a wave-based approach similar to Spotify but need to follow cursors.
+            # Because we cannot know all cursors upfront, we'll iteratively fetch waves.
+            waves: List[List[Song]] = []
+            cursors: List[str] = []
+            if next_cursor:
+                cursors.append(next_cursor)
+
+            # Limit total waves to avoid infinite loops in malformed responses
+            max_waves = 50
+            wave_size = self._playlist_concurrency * 2
+            wave_counter = 0
+            while cursors and wave_counter < max_waves:
+                current = cursors[:wave_size]
+                cursors = cursors[wave_size:]
+                tasks = [
+                    fetch_page(AM_BASE_URL + cursor) for cursor in current  # type: ignore[arg-type]
+                ]
+                results = await asyncio.gather(*tasks, return_exceptions=True)
+                for res in results:
+                    if isinstance(res, tuple):  # (songs, next)
+                        songs, nxt = res
+                        if songs:
+                            waves.append(songs)
+                        if nxt:
+                            cursors.append(nxt)
+                wave_counter += 1
+
+            for w in waves:
+                album_tracks.extend(w)
+
+            return Playlist(data, album_tracks)
 
-                    next_data: dict = await resp.json(loads=json.loads)
-                    album_tracks.extend(Song(track) for track in next_data["data"])
+    async def iter_playlist_tracks(
+        self,
+        *,
+        query: str,
+        batch_size: int = 100,
+    ) -> AsyncGenerator[List[Song], None]:
+        """Stream Apple Music playlist tracks in batches.
+
+        Parameters
+        ----------
+        query: str
+            Apple Music playlist URL.
+        batch_size: int
+            Logical grouping size for yielded batches.
+        """
+        if not self.token or datetime.utcnow() > self.expiry:
+            await self.request_token()
 
-                    _next = next_data.get("next")
-                    if _next:
-                        next_page_url = AM_BASE_URL + _next
-                    else:
-                        next_page_url = None
+        result = AM_URL_REGEX.match(query)
+        if not result or result.group("type") != "playlist":
+            raise InvalidAppleMusicURL("Provided query is not a valid Apple Music playlist URL.")
 
-            return Playlist(data, album_tracks)
+        country = result.group("country")
+        playlist_id = result.group("id")
+        request_url = AM_REQ_URL.format(country=country, type="playlist", id=playlist_id)
+        resp = await self.session.get(request_url, headers=self.headers)
+        if resp.status != 200:
+            raise AppleMusicRequestException(
+                f"Error while fetching results: {resp.status} {resp.reason}",
+            )
+        data: dict = await resp.json(loads=json.loads)
+        playlist_data = data["data"][0]
+        track_data: dict = playlist_data["relationships"]["tracks"]
+
+        first_page_tracks = [Song(track) for track in track_data["data"]]
+        for i in range(0, len(first_page_tracks), batch_size):
+            yield first_page_tracks[i : i + batch_size]
+
+        next_cursor = track_data.get("next")
+        semaphore = asyncio.Semaphore(self._playlist_concurrency)
+
+        async def fetch(cursor: str) -> tuple[List[Song], Optional[str]]:
+            url = AM_BASE_URL + cursor
+            async with semaphore:
+                r = await self.session.get(url, headers=self.headers)
+                if r.status != 200:
+                    if self._log:
+                        self._log.warning(
+                            f"Skipping Apple Music page due to {r.status} {r.reason}",
+                        )
+                    return [], None
+                pj: dict = await r.json(loads=json.loads)
+                songs = [Song(track) for track in pj.get("data", [])]
+                return songs, pj.get("next")
+
+        # Rolling waves of fetches following cursor chain
+        max_waves = 50
+        wave_size = self._playlist_concurrency * 2
+        waves = 0
+        cursors: List[str] = []
+        if next_cursor:
+            cursors.append(next_cursor)
+        while cursors and waves < max_waves:
+            current = cursors[:wave_size]
+            cursors = cursors[wave_size:]
+            results = await asyncio.gather(*[fetch(c) for c in current])
+            for songs, nxt in results:
+                if songs:
+                    for j in range(0, len(songs), batch_size):
+                        yield songs[j : j + batch_size]
+                if nxt:
+                    cursors.append(nxt)
+            waves += 1
diff --git a/pomice/enums.py b/pomice/enums.py
@@ -34,6 +34,11 @@ class SearchType(Enum):
     ytsearch = "ytsearch"
     ytmsearch = "ytmsearch"
     scsearch = "scsearch"
+    other = "other"
+
+    @classmethod
+    def _missing_(cls, value: object) -> "SearchType":  # type: ignore[override]
+        return cls.other
 
     def __str__(self) -> str:
         return self.value
@@ -68,7 +73,7 @@ class TrackType(Enum):
     OTHER = "other"
 
     @classmethod
-    def _missing_(cls, _: object) -> "TrackType":
+    def _missing_(cls, value: object) -> "TrackType":  # type: ignore[override]
         return cls.OTHER
 
     def __str__(self) -> str:
@@ -98,7 +103,7 @@ class PlaylistType(Enum):
     OTHER = "other"
 
     @classmethod
-    def _missing_(cls, _: object) -> "PlaylistType":
+    def _missing_(cls, value: object) -> "PlaylistType":  # type: ignore[override]
         return cls.OTHER
 
     def __str__(self) -> str:
@@ -213,8 +218,12 @@ class URLRegex:
 
     """
 
+    # Spotify share links can include query parameters like ?si=XXXX, a trailing slash,
+    # or an intl locale segment (e.g. /intl-en/). Broaden the regex so we still capture
+    # the type and id while ignoring extra parameters. This prevents the URL from being
+    # treated as a generic Lavalink identifier and ensures internal Spotify handling runs.
     SPOTIFY_URL = re.compile(
-        r"https?://open.spotify.com/(?P<type>album|playlist|track|artist)/(?P<id>[a-zA-Z0-9]+)",
+        r"https?://open\.spotify\.com/(?:intl-[a-zA-Z-]+/)?(?P<type>album|playlist|track|artist)/(?P<id>[a-zA-Z0-9]+)(?:/)?(?:\?.*)?$",
     )
 
     DISCORD_MP3_URL = re.compile(
@@ -235,14 +244,17 @@ class URLRegex:
         r"(?P<video>^.*?)(\?t|&start)=(?P<time>\d+)?.*",
     )
 
+    # Apple Music links sometimes append additional query parameters (e.g. &l=en, &uo=4).
+    # Allow arbitrary query parameters so valid links are captured and parsed.
     AM_URL = re.compile(
-        r"https?://music.apple.com/(?P<country>[a-zA-Z]{2})/"
-        r"(?P<type>album|playlist|song|artist)/(?P<name>.+)/(?P<id>[^?]+)",
+        r"https?://music\.apple\.com/(?P<country>[a-zA-Z]{2})/"
+        r"(?P<type>album|playlist|song|artist)/(?P<name>.+?)/(?P<id>[^/?]+?)(?:/)?(?:\?.*)?$",
     )
 
+    # Single-in-album links may also carry extra query params beyond the ?i=<trackid> token.
     AM_SINGLE_IN_ALBUM_REGEX = re.compile(
-        r"https?://music.apple.com/(?P<country>[a-zA-Z]{2})/(?P<type>album|playlist|song|artist)/"
-        r"(?P<name>.+)/(?P<id>.+)(\?i=)(?P<id2>.+)",
+        r"https?://music\.apple\.com/(?P<country>[a-zA-Z]{2})/(?P<type>album|playlist|song|artist)/"
+        r"(?P<name>.+)/(?P<id>[^/?]+)(\?i=)(?P<id2>[^&]+)(?:&.*)?$",
     )
 
     SOUNDCLOUD_URL = re.compile(
diff --git a/pomice/pool.py b/pomice/pool.py
@@ -21,7 +21,12 @@
 from discord import Client
 from discord.ext import commands
 from discord.utils import MISSING
-from websockets import client
+
+try:
+    from websockets.legacy import client  # websockets >= 10.0
+except ImportError:
+    import websockets.client as client  # websockets < 10.0  # type: ignore
+
 from websockets import exceptions
 from websockets import typing as wstype
 
@@ -303,7 +308,7 @@ async def _configure_resuming(self) -> None:
         if not self._resume_key:
             return
 
-        data = {"timeout": self._resume_timeout}
+        data: Dict[str, Union[int, str, bool]] = {"timeout": self._resume_timeout}
 
         if self._version.major == 3:
             data["resumingKey"] = self._resume_key
@@ -444,7 +449,17 @@ async def connect(self, *, reconnect: bool = False) -> Node:
         start = time.perf_counter()
 
         if not self._session:
-            self._session = aiohttp.ClientSession()
+            # Configure connection pooling for optimal concurrent request performance
+            connector = aiohttp.TCPConnector(
+                limit=100,  # Total connection limit
+                limit_per_host=30,  # Per-host connection limit
+                ttl_dns_cache=300,  # DNS cache TTL in seconds
+            )
+            timeout = aiohttp.ClientTimeout(total=30, connect=10)
+            self._session = aiohttp.ClientSession(
+                connector=connector,
+                timeout=timeout,
+            )
 
         try:
             if not reconnect:
@@ -463,7 +478,7 @@ async def connect(self, *, reconnect: bool = False) -> Node:
                         f"Version check from Node {self._identifier} successful. Returned version {version}",
                     )
 
-            self._websocket = await client.connect(
+            self._websocket = await client.connect(  # type: ignore
                 f"{self._websocket_uri}/v{self._version.major}/websocket",
                 extra_headers=self._headers,
                 ping_interval=self._heartbeat,
@@ -560,7 +575,7 @@ async def get_tracks(
         query: str,
         *,
         ctx: Optional[commands.Context] = None,
-        search_type: SearchType | None = SearchType.ytsearch,
+        search_type: Optional[SearchType] = SearchType.ytsearch,
         filters: Optional[List[Filter]] = None,
     ) -> Optional[Union[Playlist, List[Track]]]:
         """Fetches tracks from the node's REST api to parse into Lavalink.
@@ -595,7 +610,7 @@ async def get_tracks(
                         track_id=apple_music_results.id,
                         ctx=ctx,
                         track_type=TrackType.APPLE_MUSIC,
-                        search_type=search_type,
+                        search_type=search_type or SearchType.ytsearch,
                         filters=filters,
                         info={
                             "title": apple_music_results.name,
@@ -617,7 +632,7 @@ async def get_tracks(
                     track_id=track.id,
                     ctx=ctx,
                     track_type=TrackType.APPLE_MUSIC,
-                    search_type=search_type,
+                    search_type=search_type or SearchType.ytsearch,
                     filters=filters,
                     info={
                         "title": track.name,
@@ -655,7 +670,7 @@ async def get_tracks(
                         track_id=spotify_results.id,
                         ctx=ctx,
                         track_type=TrackType.SPOTIFY,
-                        search_type=search_type,
+                        search_type=search_type or SearchType.ytsearch,
                         filters=filters,
                         info={
                             "title": spotify_results.name,
@@ -677,7 +692,7 @@ async def get_tracks(
                     track_id=track.id,
                     ctx=ctx,
                     track_type=TrackType.SPOTIFY,
-                    search_type=search_type,
+                    search_type=search_type or SearchType.ytsearch,
                     filters=filters,
                     info={
                         "title": track.name,
diff --git a/pomice/spotify/client.py b/pomice/spotify/client.py

Original file line number	Diff line number	Diff line change
`@@ -20,7 +20,7 @@ class DiscordPyOutdated(Exception):`
`20`	`20`	`"using 'pip install discord.py'",`
`21`	`21`	`)`
`22`	`22`
`23`		`-__version__ = "2.9.2"`
	`23`	`+__version__ = "2.10.0"`
`24`	`24`	`__title__ = "pomice"`
`25`	`25`	`__author__ = "cloudwithax"`
`26`	`26`	`__license__ = "GPL-3.0"`