repurpose hf_token

app.py

@@ -4,6 +4,7 @@ from datetime import datetime
 
 import gradio as gr
 from gradio_client import Client
+from extract import extract
 
 functionality_shared = [
     'games',
@@ -30,19 +31,13 @@ space = 'stevez-bbc123'
 # url = f'https://{space}.hf.space/api/predict'
 url = f'https://{space}.hf.space'
 
-def extract(hf_token):
-    token = None
-    try:
-        info = json.loads(hf_token)
-        token = info['token']
-    except:
-        token = hf_token
-    print(f't = {token[:5]}')
-
 def run(hf_token, service, game, functionality, nlp_command):
     if hf_token is None or hf_token == '':
         return 'please specify hf token'
 
+    # reuse hf_token field as json string
+    token, user_name, _ = extract(hf_token)
+
     if service not in functionality_shared[1:]:
         if game is None:
             return "please specify which game"

extract.py

@@ -0,0 +1,20 @@
+import json
+
+def extract(hf_token):
+    info = {}
+    try:
+        info = json.loads(hf_token)
+    except json.decoder.JSONDecodeError:
+        return hf_token, None, None
+
+    try:
+        token = info['token']
+    except KeyError:
+        token = None
+
+    try:
+        user = info['user']
+    except KeyError:
+        user = None
+
+    return token, user, info

test_extract.py

@@ -0,0 +1,44 @@
+import pytest
+import json
+
+from extract import extract
+
+def test_pure_string():
+    inp = 'abc'
+    token, user, info = extract(inp)
+    assert token == inp
+    assert user is None
+    assert info is None
+
+def test_json_format_01():
+    inp_s = '{"a":"b"}'
+    token, user, info = extract(inp_s)
+    inp = json.loads(inp_s)
+    assert token == None
+    assert user == None
+    assert info == inp
+
+def test_json_format_02():
+    inp_s = '{"token":"b"}'
+    token, user, info = extract(inp_s)
+    inp = json.loads(inp_s)
+    assert token == inp['token']
+    assert user == None
+    assert info == inp
+
+def test_json_format_03():
+    inp_s = '{"user":"xb"}'
+    token, user, info = extract(inp_s)
+    inp = json.loads(inp_s)
+    assert token == None
+    assert user == inp['user']
+    assert info == inp
+
+def test_json_format_04():
+    inp_s = '{"user":"t1", "token":"t2"}'
+    token, user, info = extract(inp_s)
+    inp = json.loads(inp_s)
+    assert token == inp['token']
+    assert user == inp['user']
+    assert info == inp
+