From dccf09fea89eecff5feaa3a9000ecdc45e74e661 Mon Sep 17 00:00:00 2001
From: RemixDev <RemixDev64@gmail.com>
Date: Mon, 2 Aug 2021 20:42:47 +0200
Subject: [PATCH] Added errors to get_tracks_url and moved errors in separate
 file

---
 deezer/__init__.py | 21 ++++++++++++++++-----
 deezer/api.py      | 31 +++----------------------------
 deezer/errors.py   | 44 ++++++++++++++++++++++++++++++++++++++++++++
 deezer/gw.py       |  5 +----
 4 files changed, 64 insertions(+), 37 deletions(-)
 create mode 100644 deezer/errors.py

diff --git a/deezer/__init__.py b/deezer/__init__.py
index 7edeb42..518aa68 100644
--- a/deezer/__init__.py
+++ b/deezer/__init__.py
@@ -1,6 +1,8 @@
 import requests
+import json
 from deezer.gw import GW
 from deezer.api import API
+from deezer.errors import DeezerError, WrongLicense, WrongGeolocation
 
 __version__ = "1.0.4"
 
@@ -116,7 +118,8 @@ class Deezer:
                     'picture': child.get("USER_PICTURE", ""),
                     'license_token': user_data["USER"]["OPTIONS"]["license_token"],
                     'can_stream_hq': user_data["USER"]["OPTIONS"]["web_hq"] or user_data["USER"]["OPTIONS"]["mobile_hq"],
-                    'can_stream_lossless': user_data["USER"]["OPTIONS"]["web_lossless"] or user_data["USER"]["OPTIONS"]["mobile_lossless"]
+                    'can_stream_lossless': user_data["USER"]["OPTIONS"]["web_lossless"] or user_data["USER"]["OPTIONS"]["mobile_lossless"],
+                    'country': user_data["COUNTRY"]
                 })
         else:
             self.childs.append({
@@ -125,7 +128,8 @@ class Deezer:
                 'picture': user_data["USER"].get("USER_PICTURE", ""),
                 'license_token': user_data["USER"]["OPTIONS"]["license_token"],
                 'can_stream_hq': user_data["USER"]["OPTIONS"]["web_hq"] or user_data["USER"]["OPTIONS"]["mobile_hq"],
-                'can_stream_lossless': user_data["USER"]["OPTIONS"]["web_lossless"] or user_data["USER"]["OPTIONS"]["mobile_lossless"]
+                'can_stream_lossless': user_data["USER"]["OPTIONS"]["web_lossless"] or user_data["USER"]["OPTIONS"]["mobile_lossless"],
+                'country': user_data["COUNTRY"]
             })
 
     def change_account(self, child_n):
@@ -142,7 +146,9 @@ class Deezer:
         if not isinstance(track_tokens, list):
             track_tokens = [track_tokens, ]
         if not self.current_user['license_token']:
-            return []
+            return None
+        if track_format == "FLAC" and not self.current_user.can_stream_lossless or format == "MP3_320" and not self.current_user.can_stream_hq:
+            raise WrongLicense(format)
 
         try:
             request = self.session.post(
@@ -164,6 +170,11 @@ class Deezer:
         except requests.exceptions.HTTPError:
             return None
 
-        if response.get('data') and response['data'][0].get('media'):
-            return response['data'][0]['media'][0]['sources'][0]['url']
+        if len(response.get('data', [])):
+            if response['data'][0]['errors']:
+                if response['data'][0]['errors'][0]['code'] == 2002:
+                    raise WrongGeolocation(self.current_user.country)
+                raise DeezerError(json.dumps(response))
+            if response['data'][0]['media']:
+                return response['data'][0]['media'][0]['sources'][0]['url']
         return None
diff --git a/deezer/api.py b/deezer/api.py
index e31d0df..19d67d5 100644
--- a/deezer/api.py
+++ b/deezer/api.py
@@ -2,6 +2,9 @@ import requests
 from time import sleep
 
 import json
+from deezer.errors import ItemsLimitExceededException, PermissionException, InvalidTokenException, \
+WrongParameterException, MissingParameterException, InvalidQueryException, DataException, \
+IndividualAccountChangedNotAllowedException, APIError
 
 class SearchOrder():
     """Possible values for order parameter in search"""
@@ -289,31 +292,3 @@ class API:
             resp = self.advanced_search(artist=artist, track=track[:track.find(" - ")], limit=1)
             if len(resp['data']) > 0: return resp['data'][0]['id']
         return "0"
-
-class APIError(Exception):
-    """Base class for Deezer exceptions"""
-    pass
-
-class ItemsLimitExceededException(APIError):
-    pass
-
-class PermissionException(APIError):
-    pass
-
-class InvalidTokenException(APIError):
-    pass
-
-class WrongParameterException(APIError):
-    pass
-
-class MissingParameterException(APIError):
-    pass
-
-class InvalidQueryException(APIError):
-    pass
-
-class DataException(APIError):
-    pass
-
-class IndividualAccountChangedNotAllowedException(APIError):
-    pass
diff --git a/deezer/errors.py b/deezer/errors.py
new file mode 100644
index 0000000..a55fd70
--- /dev/null
+++ b/deezer/errors.py
@@ -0,0 +1,44 @@
+class DeezerError(Exception):
+    """Base class for Deezer exceptions"""
+
+class WrongLicense(DeezerError):
+    def __init__(self, track_format):
+        super().__init__()
+        self.message = f"Your account doesn't have the license to stream {track_format}"
+        self.format = track_format
+
+class WrongGeolocation(DeezerError):
+    def __init__(self, country):
+        super().__init__()
+        self.message = f"The track you requested can't be streamed in country {country}"
+        self.country = country
+
+class APIError(DeezerError):
+    """Base class for Deezer api exceptions"""
+
+class ItemsLimitExceededException(APIError):
+    pass
+
+class PermissionException(APIError):
+    pass
+
+class InvalidTokenException(APIError):
+    pass
+
+class WrongParameterException(APIError):
+    pass
+
+class MissingParameterException(APIError):
+    pass
+
+class InvalidQueryException(APIError):
+    pass
+
+class DataException(APIError):
+    pass
+
+class IndividualAccountChangedNotAllowedException(APIError):
+    pass
+
+class GWAPIError(DeezerError):
+    """Base class for Deezer gw api exceptions"""
diff --git a/deezer/gw.py b/deezer/gw.py
index edd9d47..afbebde 100644
--- a/deezer/gw.py
+++ b/deezer/gw.py
@@ -3,6 +3,7 @@ from time import sleep
 
 import json
 from deezer.utils import map_artist_album, map_user_track, map_user_artist, map_user_album, map_user_playlist
+from deezer.errors import GWAPIError
 
 class LyricsStatus():
     """Explicit Content Lyrics"""
@@ -374,7 +375,3 @@ class GW:
         for track in data:
             result.append(map_user_track(track))
         return result
-
-class GWAPIError(Exception):
-    """Base class for Deezer exceptions"""
-    pass