From a783edd3ac619fc678bda4685de76aa981eceab2 Mon Sep 17 00:00:00 2001
From: Bradlee Speice <bradlee@speice.io>
Date: Sat, 3 Jul 2021 22:58:51 -0400
Subject: [PATCH] More type-friendly paging API

---
 spotify_actions/search.py | 11 ++++-------
 spotify_actions/util.py   | 13 ++++++-------
 2 files changed, 10 insertions(+), 14 deletions(-)

diff --git a/spotify_actions/search.py b/spotify_actions/search.py
index a27e1e0..48d1457 100644
--- a/spotify_actions/search.py
+++ b/spotify_actions/search.py
@@ -1,9 +1,9 @@
 """
 Utility methods for the Spotify query API
 """
-from typing import Any, Dict, Iterable, cast
+from typing import Iterable
 
-from spotify_model.album import SearchAlbum
+from spotify_model import Paging, SearchAlbum
 from spotipy import Spotify
 
 from .util import exhaust
@@ -12,11 +12,8 @@ from .util import exhaust
 def search_album(client: Spotify, search_str: str) -> Iterable[SearchAlbum]:
     "Display albums from a search string"
 
-    def _search(limit: int, offset: int) -> Dict[str, Any]:
-        return cast(
-            Dict[str, Any],
-            client.search(search_str, limit=limit, offset=offset, type="album")["albums"],
-        )
+    def _search(limit: int, offset: int) -> Paging:
+        return Paging(**client.search(search_str, limit=limit, offset=offset, type="album")["albums"])
 
     for item in exhaust(_search):
         yield SearchAlbum(**item)
diff --git a/spotify_actions/util.py b/spotify_actions/util.py
index 443f81b..e9b88e8 100644
--- a/spotify_actions/util.py
+++ b/spotify_actions/util.py
@@ -6,6 +6,7 @@ from pathlib import Path
 from typing import Any, Dict, Iterable, Protocol
 
 import yaml
+from spotify_model import Paging
 from spotipy import Spotify, SpotifyClientCredentials
 
 DEFAULT_LIMIT = 50
@@ -23,20 +24,18 @@ def read_credentials(path: Path) -> Spotify:
 class Paginated(Protocol):
     "Protocol definition for functions that will be provided to the `exhaust` handler"
 
-    def __call__(self, limit: int, offset: int) -> Dict[str, Any]:
+    def __call__(self, limit: int, offset: int) -> Paging:
         ...
 
 
 def exhaust(function: Paginated, limit: int = DEFAULT_LIMIT) -> Iterable[Dict[str, Any]]:
     "Exhaust a function that returns a pagination object"
     response = function(limit=limit, offset=0)
-    total = response["total"]
-    limit = response["limit"]
 
-    for item in response["items"]:
+    for item in response.items:
         yield item
 
-    for i in range(1, ceil(total / limit)):
-        response = function(limit=limit, offset=limit * i)
-        for item in response["items"]:
+    for i in range(1, ceil(response.total / limit)):
+        response = function(limit=response.limit, offset=response.limit * i)
+        for item in response.items:
             yield item