From 1ee7b65d71011893549751946570804860778b53 Mon Sep 17 00:00:00 2001 From: Crussader <75786691+Crussader@users.noreply.github.com> Date: Sun, 27 Mar 2022 20:57:36 +0400 Subject: [PATCH] Update client.py --- pomice/spotify/client.py | 143 ++++++++++++++++++++++++++++----------- 1 file changed, 103 insertions(+), 40 deletions(-) diff --git a/pomice/spotify/client.py b/pomice/spotify/client.py index cc3cf21..1baf87c 100644 --- a/pomice/spotify/client.py +++ b/pomice/spotify/client.py @@ -1,18 +1,27 @@ +import asyncio import re +import itertools import time from base64 import b64encode import aiohttp +try: + import orjson + json = orjson +except ImportError: + import json + from .album import Album from .exceptions import InvalidSpotifyURL, SpotifyRequestException from .playlist import Playlist from .track import Track +from .artist import Artist GRANT_URL = "https://accounts.spotify.com/api/token" REQUEST_URL = "https://api.spotify.com/v1/{type}s/{id}" SPOTIFY_URL_REGEX = re.compile( - r"https?://open.spotify.com/(?Palbum|playlist|track)/(?P[a-zA-Z0-9]+)" + r"https?://open.spotify.com/(?Palbum|playlist|track|artist)/(?P[a-zA-Z0-9]+)" ) @@ -22,7 +31,12 @@ class Client: for any Spotify URL you throw at it. """ - def __init__(self, client_id: str, client_secret: str) -> None: + def __init__( + self, + client_id: str, + client_secret: str, + ) -> None: + self._client_id = client_id self._client_secret = client_secret @@ -30,8 +44,10 @@ class Client: self._bearer_token: str = None self._expiry = 0 - self._auth_token = b64encode(f"{self._client_id}:{self._client_secret}".encode()) - self._grant_headers = {"Authorization": f"Basic {self._auth_token.decode()}"} + self._auth_token = b64encode( + f"{self._client_id}:{self._client_secret}".encode()) + self._grant_headers = { + "Authorization": f"Basic {self._auth_token.decode()}"} self._bearer_headers = None async def _fetch_bearer_token(self) -> None: @@ -43,15 +59,69 @@ class Client: f"Error fetching bearer token: {resp.status} {resp.reason}" ) - data: dict = await resp.json() + data: dict = await resp.json(loads=json.loads) self._bearer_token = data["access_token"] self._expiry = time.time() + (int(data["expires_in"]) - 10) - self._bearer_headers = {"Authorization": f"Bearer {self._bearer_token}"} - - + self._bearer_headers = { + "Authorization": f"Bearer {self._bearer_token}"} - async def search(self, *, query: str): + async def _load(self, type_: str, request_url: str, data: dict, tracks: list, other: dict = None, full: bool = False): + + async def _fetch_async(count, url): + async with self.session.get(url, headers=self._bearer_headers) as resp: + if resp.status != 200: + raise SpotifyRequestException( + f"Error while fetching results: {resp.status} {resp.reason}" + ) + + next_data: dict = await resp.json(loads=json.loads) + inner = [count] + inner.extend(Track(track.get('track') or track) + for track in next_data['items']) + tracks.append(inner) + + processes = [] + + if type_ == 'playlist': + if not len(tracks): + raise SpotifyRequestException( + "This playlist is empty and therefore cannot be queued.") + + urls = [f"{request_url}/tracks?offset={off}" + for off in range(100, data['tracks']['total']+100, 100)] + + cls = Playlist(data, []) + else: + urls = [REQUEST_URL.format(type="album", id=album['id'])+'/tracks' + for album in (data['items'] if not other else other['items'])] + + cls = Artist(data, []) + + + processes = [_fetch_async(url, count) + for url, count in enumerate(urls[:9] if not full else urls, start=1)] + + try: + await asyncio.gather(*processes) + except TimeoutError: + # dont know why it TimeoutError simply happens + # but if it does then manually waiting for each func + [await func for func in processes] + + # tracks are jumbled for huge playlists so we use this to sort it out + tracks.sort(key=lambda i: False if isinstance(i, Track) else i[0]) + tracks = [track for track in itertools.chain( + *tracks) if not isinstance(track, int)] + + + cls.tracks = tracks + if cls.total_tracks == 0: + cls.total_tracks = len(tracks) + + return cls + + async def search(self, *, query: str, raw: bool = False): if not self._bearer_token or time.time() >= self._expiry: await self._fetch_bearer_token() @@ -70,41 +140,34 @@ class Client: f"Error while fetching results: {resp.status} {resp.reason}" ) - data: dict = await resp.json() + data = await resp.json(loads=json.loads) - if spotify_type == "track": + if raw: + return data + + if spotify_type == 'track': return Track(data) - elif spotify_type == "album": + elif spotify_type == 'album': return Album(data) else: + tracks = [] + if spotify_type == 'playlist': + inner = [0] + inner.extend(Track(track['track']) + for track in data['tracks']['items'] if track['track']) + tracks.append(inner) - tracks = [ - Track(track["track"]) - for track in data["tracks"]["items"] if track["track"] is not None - ] + other = None # is for artist data + else: + if not data.get('items'): + async with self.session.get(f"{request_url}/albums", headers=self._bearer_headers) as resp: + if resp.status != 200: + raise SpotifyRequestException( + f"Error while fetching results: {resp.status} {resp.reason}" + ) - if not len(tracks): - raise SpotifyRequestException("This playlist is empty and therefore cannot be queued.") - - next_page_url = data["tracks"]["next"] - - while next_page_url is not None: - async with self.session.get(next_page_url, headers=self._bearer_headers) as resp: - if resp.status != 200: - raise SpotifyRequestException( - f"Error while fetching results: {resp.status} {resp.reason}" - ) - - next_data: dict = await resp.json() - - tracks += [ - Track(track["track"]) - for track in next_data["items"] if track["track"] is not None - ] - next_page_url = next_data["next"] - - return Playlist(data, tracks) + other = await resp.json(loads=json.loads) + else: + other = None - async def close(self): - await self.session.close() - + return await self._load(spotify_type, request_url, data, tracks, other)