Move the invalidity token management system outsite of the module

2 files changed, 40 insertions, 17 deletions

diff --git a/README.md b/README.md
index 023c872..8175cd4 100644
--- a/README.md
+++ b/README.md
@@ -54,16 +54,30 @@ token_endpoint = "https://demo.funkwhale.audio/api/v1/oauth/token/"
 # Save the OAuth2 infos in file, this is ugly yup
 token_filename = 'oauth_token.data'
 
+# You need to handle the invalidity of a token and ask for a new one
+def _ask_new_auth(funkwhale, message):
+  print(message)
+  authorization_code = input("Enter response code: ")
+  funkwhale.client._set_token(authorization_code)
+
 funkwhale = Funkwhale(client_name, redirect_uri, client_id, client_secret,
                 scopes, username, password, domain, authorization_endpoint,
                 token_endpoint, token_filename)
 
-artists = funkwhale.artists()
+try:
+  artists = funkwhale.artists()
+except InvalidTokenError as e
+  _ask_new_auth(funkwhale, e.message)
+  artists = funkwhale.artists()
 
 print(artists)
 
 # In case you ask, their is an example for downloading a song
-r = funkwhale.listen("f9d02c64-bafa-43cb-8e1e-fa612e7c5dab")
+try:
+  r = funkwhale.listen("f9d02c64-bafa-43cb-8e1e-fa612e7c5dab")
+except InvalidTokenError as e:
+  _ask_new_auth(funkwhale, e.message)
+  r = funkwhale.rate_limit()
 
 with open('/tmp/test.mp3', 'wb') as f:
     for chunk in r.iter_content(chunk_size=8192):
diff --git a/pyfunkwhale/client.py b/pyfunkwhale/client.py
index 53e7a66..a2f5b57 100644
--- a/pyfunkwhale/client.py
+++ b/pyfunkwhale/client.py
@@ -10,6 +10,14 @@ from requests.models import Response
 
 from pyfunkwhale.utils import read_file, write_file
 
+class InvalidTokenError(Exception):
+
+    def __init__(self, client):
+        self.message = client.authorization_url
+
+    def __str__(self):
+        return self.message
+
 
 class Client(object):
 
@@ -39,26 +47,27 @@ class Client(object):
             authorization_url(
                             self.authorization_endpoint)
 
+    def _connect(self):
+        """
+        Use saved token or ask a new one to the API.
+        """
         try:
-            self.token = json.loads(read_file(token_filename))
+            self.token = json.loads(read_file(self.token_filename))
             self._refresh_token()
         except FileNotFoundError:
-            self._get_token()
-            write_file(token_filename, self.token)
+            raise InvalidTokenError(self)
 
-    def _get_token(self):
+    def _set_token(self, authorization_code):
         """
-        Ask the user to go on the authorization page of the funkwhale instance
-        and then wait the response code for fetch the token generated.
+        Use the authorization_code to fetch the new token.
         """
-        print("For authorizate this app go to:\n{}".format(
-                self.authorization_url))
-        self.authorization_code = input("Enter response code: ")
+        self.authorization_code = authorization_code
 
         self.token = self.oauth_client.fetch_token(
                                 self.token_endpoint,
                                 code=self.authorization_code,
                                 client_secret=self.client_secret)
+        write_file(self.token_filename, self.token)
 
     def _refresh_token(self):
         """
@@ -110,23 +119,23 @@ class Client(object):
         ------
         requests.exceptions.HTTPError
             If their is an error during requesting the API.
+        pyfunkwhale.client.InvalidTokenError
+            If current token is invalid
         """
         self._refresh_token()
         if headers is None:
             headers = {'Authorization': self.token['token_type'] + ' ' +
                        self.token['access_token']}
 
-        call = getattr(self.oauth_client, method)
+        _call = getattr(self.oauth_client, method)
 
         endpoint = re.sub(r'^\/', '', endpoint)
 
-        r = call(self.domain + '/api/v1/' + endpoint, headers=headers,
-                 params=params, data=data)
+        r = _call(self.domain + '/api/v1/' + endpoint, headers=headers,
+                params=params, data=data)
 
         if r.status_code == 401:
-            self._force_refresh_token()
-            r = call(self.domain + '/api/v1/' + endpoint, headers=headers,
-                     params=params, data=data)
+            raise InvalidTokenError(self)
 
         r.raise_for_status()