Added batching for max filter size.

Author: semohr

beetsplug/tidal/api.py

@@ -1,9 +1,10 @@
 from __future__ import annotations
 
 import urllib.parse
+from itertools import batched, zip_longest
 from pathlib import Path
 from time import sleep
-from typing import TYPE_CHECKING, Any, TypeVar
+from typing import TYPE_CHECKING, Any, TypeVar, overload
 
 import requests
 
@@ -26,6 +27,7 @@
 
 
 API_BASE = "https://openapi.tidal.com/v2"
+MAX_FILTER_SIZE = 20
 
 
 class TidalSession(TimeoutAndRetrySession):
@@ -129,7 +131,7 @@ def get_paginated(
                 **kwargs,
             )
             page_doc = res.json()
-            doc = self._merge_multiresource_pagination(doc, page_doc)
+            doc = self.merge_multiresource_pagination(doc, page_doc)
 
         # Dedupe include
         doc["included"] = list(
@@ -140,8 +142,8 @@ def get_paginated(
         )
         return doc
 
-    def _merge_multiresource_pagination(
-        self,
+    @staticmethod
+    def merge_multiresource_pagination(
         a: Document[list[T]],
         b: Document[list[T]],
     ) -> Document[list[T]]:
@@ -227,63 +229,120 @@ def search_results(
             params=params,
         ).json()
 
+    @overload
+    def get_tracks(
+        self,
+        *,
+        ids: list[str] | str,
+        include: list[str] | str | None = None,
+        country_code: str = "US",
+    ) -> TrackDocument: ...
+
+    @overload
+    def get_tracks(
+        self,
+        *,
+        isrcs: list[str] | str,
+        include: list[str] | str | None = None,
+        country_code: str = "US",
+    ) -> TrackDocument: ...
+
     def get_tracks(
         self,
-        # filters
         ids: list[str] | str | None = None,
         isrcs: list[str] | str | None = None,
-        *,
         include: list[str] | str | None = None,
         country_code: str = "US",
     ) -> TrackDocument:
         """Fetch tracks resolving pagination and included items.
 
-        Should only ever be called with 20 items as
-        tidal does not support more per requests. This does not mean more than
-        20 cant be returned.
-
         https://tidal-music.github.io/tidal-api-reference/#/tracks/get_tracks
         """
-        params: dict[str, str | list[str]] = {}
-        if country_code:
-            params["countryCode"] = country_code
-        if ids:
-            params["filter[id]"] = ids
-        if isrcs:
-            params["filter[isrc]"] = isrcs
-
-        return self.session.get_paginated(
-            f"{API_BASE}/tracks",
-            include,
-            params=params,
-        )
+        ids = [ids] if isinstance(ids, str) else ids or []
+        isrcs = [isrcs] if isinstance(isrcs, str) else isrcs or []
+
+        # Tidal allows at max 20 filters per request. This needs a bit of extra
+        # logic sadly.
+        doc: TrackDocument = {
+            "data": [],
+            "included": [],
+        }
+        for id_batch, isrc_batch in zip_longest(
+            batched(ids, MAX_FILTER_SIZE),
+            batched(isrcs, MAX_FILTER_SIZE),
+            fillvalue=(),
+        ):
+            params: dict[str, Any] = {"countryCode": country_code}
+            if id_batch:
+                params["filter[id]"] = id_batch
+            if isrc_batch:
+                params["filter[isrc]"] = isrc_batch
+
+            doc = self.session.merge_multiresource_pagination(
+                doc,
+                self.session.get_paginated(
+                    f"{API_BASE}/tracks", include, params=params
+                ),
+            )
+
+        return doc
+
+    @overload
+    def get_albums(
+        self,
+        *,
+        ids: list[str] | str,
+        include: list[str] | str | None = None,
+        country_code: str = "US",
+    ) -> AlbumDocument: ...
+
+    @overload
+    def get_albums(
+        self,
+        *,
+        barcode_ids: list[str] | str,
+        include: list[str] | str | None = None,
+        country_code: str = "US",
+    ) -> AlbumDocument: ...
 
     def get_albums(
         self,
-        # filters
         ids: list[str] | str | None = None,
         barcode_ids: list[str] | str | None = None,
-        *,
         include: list[str] | str | None = None,
         country_code: str = "US",
     ) -> AlbumDocument:
         """Fetch Albums resolving pagination and included items.
 
-        Should only ever be called with 20 items as tidal does not support more per
-        requests. This does not mean more than 20 cant be returned.
-
         https://tidal-music.github.io/tidal-api-reference/#/albums/get_albums
         """
-        params: dict[str, str | list[str]] = {}
-        if country_code:
-            params["countryCode"] = country_code
-        if ids:
-            params["filter[id]"] = ids
-        if barcode_ids:
-            params["filter[barcodeId]"] = barcode_ids
-
-        return self.session.get_paginated(
-            f"{API_BASE}/albums",
-            include,
-            params=params,
+        ids = [ids] if isinstance(ids, str) else ids or []
+        barcode_ids = (
+            [barcode_ids] if isinstance(barcode_ids, str) else barcode_ids or []
         )
+
+        # Tidal allows at max 20 filters per request. This needs a bit of extra
+        # logic sadly.
+        doc: AlbumDocument = {
+            "data": [],
+            "included": [],
+        }
+        for id_batch, barcode_batch in zip_longest(
+            batched(ids, MAX_FILTER_SIZE),
+            batched(barcode_ids, MAX_FILTER_SIZE),
+            fillvalue=(),
+        ):
+            params: dict[str, Any] = {"countryCode": country_code}
+            if id_batch:
+                params["filter[id]"] = id_batch
+            if barcode_batch:
+                params["filter[barcodeId]"] = barcode_batch
+
+            doc = self.session.merge_multiresource_pagination(
+                doc,
+                self.session.get_paginated(
+                    f"{API_BASE}/albums", include, params=params
+                ),
+            )
+
+        return doc