Repository: alexgolec/tda-api

Issue:
Current token format is not portable
Pickle is not a portable data format, we should use JSON which is generally the standard for storing oauth credentials across multiple libraries.

Comments:

Patch:
diff --git a/tda/auth.py b/tda/auth.py
index ff8d2cf..fec6b1e 100644
--- a/tda/auth.py
+++ b/tda/auth.py
@@ -3,6 +3,7 @@
 
 from requests_oauthlib import OAuth2Session
 
+import json
 import logging
 import pickle
 import time
@@ -19,11 +20,28 @@ def __token_updater(token_path):
     def update_token(t):
         get_logger().info('Updating token to file {}'.format(token_path))
 
-        with open(token_path, 'wb') as f:
-            pickle.dump(t, f)
+        with open(token_path, 'w') as f:
+            json.dump(t, f)
     return update_token
 
 
+def __token_loader(token_path):
+    def load_token():
+        get_logger().info('Loading token from file {}'.format(token_path))
+
+        with open(token_path, 'rb') as f:
+            token_data = f.read()
+            try:
+                return json.loads(token_data.decode())
+            except ValueError:
+                get_logger().warning(
+                    "Unable to load JSON token from file {}, falling back to pickle"\
+                    .format(token_path)
+                )
+                return pickle.loads(token_data)
+    return load_token
+
+
 def __normalize_api_key(api_key):
     api_key_suffix = '@AMER.OAUTHAP'
 
@@ -51,9 +69,8 @@ def client_from_token_file(token_path, api_key):
     :param api_key: Your TD Ameritrade application's API key, also known as the
                     client ID.
     '''
-    def load():
-        with open(token_path, 'rb') as f:
-            return pickle.load(f)
+    
+    load = __token_loader(token_path)
 
     return client_from_access_functions(
             api_key, load, __token_updater(token_path))


Test patch:
diff --git a/tests/auth_test.py b/tests/auth_test.py
index adbd092..8b165b0 100644
--- a/tests/auth_test.py
+++ b/tests/auth_test.py
@@ -3,6 +3,7 @@ from tests.test_utils import no_duplicates
 from unittest.mock import patch, ANY, MagicMock
 from unittest.mock import ANY as _
 
+import json
 import os
 import pickle
 import tempfile
@@ -17,11 +18,14 @@ class ClientFromTokenFileTest(unittest.TestCase):
     def setUp(self):
         self.tmp_dir = tempfile.TemporaryDirectory()
         self.pickle_path = os.path.join(self.tmp_dir.name, 'token.pickle')
+        self.json_path = os.path.join(self.tmp_dir.name, 'token.json')
         self.token = {'token': 'yes'}
 
     def write_token(self):
         with open(self.pickle_path, 'wb') as f:
             pickle.dump(self.token, f)
+        with open(self.json_path, 'w') as f:
+            json.dump(self.token, f)
 
     @no_duplicates
     def test_no_such_file(self):
@@ -31,7 +35,7 @@ class ClientFromTokenFileTest(unittest.TestCase):
     @no_duplicates
     @patch('tda.auth.Client')
     @patch('tda.auth.OAuth2Session')
-    def test_file_exists(self, session, client):
+    def test_pickle_loads(self, session, client):
         self.write_token()
 
         client.return_value = 'returned client'
@@ -46,13 +50,31 @@ class ClientFromTokenFileTest(unittest.TestCase):
             auto_refresh_kwargs=_,
             token_updater=_)
 
+    @no_duplicates
+    @patch('tda.auth.Client')
+    @patch('tda.auth.OAuth2Session')
+    def test_json_loads(self, session, client):
+        self.write_token()
+
+        client.return_value = 'returned client'
+
+        self.assertEqual('returned client',
+                         auth.client_from_token_file(self.json_path, API_KEY))
+        client.assert_called_once_with(API_KEY, _)
+        session.assert_called_once_with(
+            API_KEY,
+            token=self.token,
+            auto_refresh_url=_,
+            auto_refresh_kwargs=_,
+            token_updater=_)
+
     @no_duplicates
     @patch('tda.auth.Client')
     @patch('tda.auth.OAuth2Session')
     def test_token_updater_updates_token(self, session, client):
         self.write_token()
 
-        auth.client_from_token_file(self.pickle_path, API_KEY)
+        auth.client_from_token_file(self.json_path, API_KEY)
         session.assert_called_once()
 
         session_call = session.mock_calls[0]
@@ -60,8 +82,8 @@ class ClientFromTokenFileTest(unittest.TestCase):
 
         updated_token = {'updated': 'token'}
         token_updater(updated_token)
-        with open(self.pickle_path, 'rb') as f:
-            self.assertEqual(pickle.load(f), updated_token)
+        with open(self.json_path, 'r') as f:
+            self.assertEqual(json.load(f), updated_token)
 
 
     @no_duplicates
@@ -149,7 +171,7 @@ class ClientFromLoginFlow(unittest.TestCase):
 
     def setUp(self):
         self.tmp_dir = tempfile.TemporaryDirectory()
-        self.pickle_path = os.path.join(self.tmp_dir.name, 'token.pickle')
+        self.json_path = os.path.join(self.tmp_dir.name, 'token.json')
         self.token = {'token': 'yes'}
 
     @no_duplicates
@@ -171,11 +193,11 @@ class ClientFromLoginFlow(unittest.TestCase):
         self.assertEqual('returned client',
                          auth.client_from_login_flow(
                              webdriver, API_KEY, REDIRECT_URL,
-                             self.pickle_path,
+                             self.json_path,
                              redirect_wait_time_seconds=0.0))
 
-        with open(self.pickle_path, 'rb') as f:
-            self.assertEqual(self.token, pickle.load(f))
+        with open(self.json_path, 'r') as f:
+            self.assertEqual(self.token, json.load(f))
 
     @no_duplicates
     @patch('tda.auth.Client')
@@ -198,11 +220,11 @@ class ClientFromLoginFlow(unittest.TestCase):
         self.assertEqual('returned client',
                          auth.client_from_login_flow(
                              webdriver, API_KEY, redirect_url,
-                             self.pickle_path,
+                             self.json_path,
                              redirect_wait_time_seconds=0.0))
 
-        with open(self.pickle_path, 'rb') as f:
-            self.assertEqual(self.token, pickle.load(f))
+        with open(self.json_path, 'r') as f:
+            self.assertEqual(self.token, json.load(f))
 
     @no_duplicates
     @patch('tda.auth.Client')
@@ -227,11 +249,11 @@ class ClientFromLoginFlow(unittest.TestCase):
         self.assertEqual('returned client',
                          auth.client_from_login_flow(
                              webdriver, API_KEY, redirect_url,
-                             self.pickle_path,
+                             self.json_path,
                              redirect_wait_time_seconds=0.0))
 
-        with open(self.pickle_path, 'rb') as f:
-            self.assertEqual(self.token, pickle.load(f))
+        with open(self.json_path, 'r') as f:
+            self.assertEqual(self.token, json.load(f))
 
     @no_duplicates
     @patch('tda.auth.Client')
@@ -252,7 +274,7 @@ class ClientFromLoginFlow(unittest.TestCase):
         self.assertEqual('returned client',
                          auth.client_from_login_flow(
                              webdriver, 'API_KEY', REDIRECT_URL,
-                             self.pickle_path,
+                             self.json_path,
                              redirect_wait_time_seconds=0.0))
 
         self.assertEqual(
@@ -280,7 +302,7 @@ class ClientFromLoginFlow(unittest.TestCase):
                 'timed out waiting for redirect'):
             auth.client_from_login_flow(
                     webdriver, API_KEY, redirect_url,
-                    self.pickle_path,
+                    self.json_path,
                     redirect_wait_time_seconds=0.0)
 
 
@@ -288,12 +310,12 @@ class EasyClientTest(unittest.TestCase):
 
     def setUp(self):
         self.tmp_dir = tempfile.TemporaryDirectory()
-        self.pickle_path = os.path.join(self.tmp_dir.name, 'token.pickle')
+        self.json_path = os.path.join(self.tmp_dir.name, 'token.json')
         self.token = {'token': 'yes'}
 
     def write_token(self):
-        with open(self.pickle_path, 'wb') as f:
-            pickle.dump(self.token, f)
+        with open(self.json_path, 'w') as f:
+            json.dump(self.token, f)
 
     @no_duplicates
     @patch('tda.auth.client_from_token_file')
@@ -302,7 +324,7 @@ class EasyClientTest(unittest.TestCase):
         client_from_token_file.side_effect = FileNotFoundError()
 
         with self.assertRaises(FileNotFoundError):
-            auth.easy_client(API_KEY, REDIRECT_URL, self.pickle_path)
+            auth.easy_client(API_KEY, REDIRECT_URL, self.json_path)
 
     @no_duplicates
     @patch('tda.auth.client_from_token_file')
@@ -311,7 +333,7 @@ class EasyClientTest(unittest.TestCase):
         client_from_token_file.return_value = self.token
 
         self.assertEquals(self.token,
-                          auth.easy_client(API_KEY, REDIRECT_URL, self.pickle_path))
+                          auth.easy_client(API_KEY, REDIRECT_URL, self.json_path))
 
     @no_duplicates
     @patch('tda.auth.client_from_login_flow')
@@ -327,7 +349,7 @@ class EasyClientTest(unittest.TestCase):
 
         self.assertEquals('returned client',
                           auth.easy_client(
-                              API_KEY, REDIRECT_URL, self.pickle_path,
+                              API_KEY, REDIRECT_URL, self.json_path,
                               webdriver_func=webdriver_func))
 
         webdriver_func.assert_called_once()