Merge pull request #676 from chripede/django-jsonpickle

Use jsonpickle in django contrib

6 files changed, 24 insertions, 6 deletions

diff --git a/oauth2client/contrib/devshell.py b/oauth2client/contrib/devshell.py
index c1906eb..691765f 100644
--- a/oauth2client/contrib/devshell.py
+++ b/oauth2client/contrib/devshell.py
@@ -37,6 +37,7 @@ class CommunicationError(Error):
 class NoDevshellServer(Error):
     """Error when no Developer Shell server can be contacted."""
 
+
 # The request for credential information to the Developer Shell client socket
 # is always an empty PBLite-formatted JSON object, so just define it as a
 # constant.
diff --git a/oauth2client/contrib/django_util/models.py b/oauth2client/contrib/django_util/models.py
index 87e1da7..37cc697 100644
--- a/oauth2client/contrib/django_util/models.py
+++ b/oauth2client/contrib/django_util/models.py
@@ -19,6 +19,7 @@ import pickle
 
 from django.db import models
 from django.utils import encoding
+import jsonpickle
 
 import oauth2client
 
@@ -48,7 +49,12 @@ class CredentialsField(models.Field):
         elif isinstance(value, oauth2client.client.Credentials):
             return value
         else:
-            return pickle.loads(base64.b64decode(encoding.smart_bytes(value)))
+            try:
+                return jsonpickle.decode(
+                    base64.b64decode(encoding.smart_bytes(value)).decode())
+            except ValueError:
+                return pickle.loads(
+                    base64.b64decode(encoding.smart_bytes(value)))
 
     def get_prep_value(self, value):
         """Overrides ``models.Field`` method. This is used to convert
@@ -58,7 +64,8 @@ class CredentialsField(models.Field):
         if value is None:
             return None
         else:
-            return encoding.smart_text(base64.b64encode(pickle.dumps(value)))
+            return encoding.smart_text(
+                base64.b64encode(jsonpickle.encode(value).encode()))
 
     def value_to_string(self, obj):
         """Convert the field value from the provided model to a string.
diff --git a/oauth2client/tools.py b/oauth2client/tools.py
index b882429..5166993 100644
--- a/oauth2client/tools.py
+++ b/oauth2client/tools.py
@@ -92,6 +92,7 @@ def _CreateArgumentParser():
         help='Set the logging level of detail.')
     return parser
 
+
 # argparser is an ArgumentParser that contains command-line options expected
 # by tools.run(). Pass it in as part of the 'parents' argument to your own
 # ArgumentParser.
diff --git a/tests/contrib/django_util/test_django_models.py b/tests/contrib/django_util/test_django_models.py
index f87250d..da54965 100644
--- a/tests/contrib/django_util/test_django_models.py
+++ b/tests/contrib/django_util/test_django_models.py
@@ -21,6 +21,8 @@ import base64
 import pickle
 import unittest
 
+import jsonpickle
+
 from oauth2client import _helpers
 from oauth2client import client
 from oauth2client.contrib.django_util import models
@@ -36,6 +38,8 @@ class TestCredentialsField(unittest.TestCase):
         self.credentials = client.Credentials()
         self.pickle_str = _helpers._from_bytes(
             base64.b64encode(pickle.dumps(self.credentials)))
+        self.jsonpickle_str = _helpers._from_bytes(
+            base64.b64encode(jsonpickle.encode(self.credentials).encode()))
 
     def test_field_is_text(self):
         self.assertEqual(self.field.get_internal_type(), 'BinaryField')
@@ -44,6 +48,10 @@ class TestCredentialsField(unittest.TestCase):
         self.assertIsInstance(
             self.field.to_python(self.pickle_str), client.Credentials)
 
+    def test_field_jsonunpickled(self):
+        self.assertIsInstance(
+            self.field.to_python(self.jsonpickle_str), client.Credentials)
+
     def test_field_already_unpickled(self):
         self.assertIsInstance(
             self.field.to_python(self.credentials), client.Credentials)
@@ -62,12 +70,12 @@ class TestCredentialsField(unittest.TestCase):
     def test_field_pickled(self):
         prep_value = self.field.get_db_prep_value(self.credentials,
                                                   connection=None)
-        self.assertEqual(prep_value, self.pickle_str)
+        self.assertEqual(prep_value, self.jsonpickle_str)
 
     def test_field_value_to_string(self):
         self.fake_model.credentials = self.credentials
         value_str = self.fake_model_field.value_to_string(self.fake_model)
-        self.assertEqual(value_str, self.pickle_str)
+        self.assertEqual(value_str, self.jsonpickle_str)
 
     def test_field_value_to_string_none(self):
         self.fake_model.credentials = None
diff --git a/tests/contrib/test_devshell.py b/tests/contrib/test_devshell.py
index d5aea32..1346080 100644
--- a/tests/contrib/test_devshell.py
+++ b/tests/contrib/test_devshell.py
@@ -160,8 +160,8 @@ class _AuthReferenceServer(threading.Thread):
                     s.recv(to_read, socket.MSG_WAITALL))
             if resp_buffer != devshell.CREDENTIAL_INFO_REQUEST_JSON:
                 self.bad_request = True
-            l = len(self.response)
-            s.sendall('{0}\n{1}'.format(l, self.response).encode())
+            response_len = len(self.response)
+            s.sendall('{0}\n{1}'.format(response_len, self.response).encode())
         finally:
             # Will fail if s is None, but these tests never encounter
             # that scenario.
diff --git a/tests/test_service_account.py b/tests/test_service_account.py
index 8d93936..6756d49 100644
--- a/tests/test_service_account.py
+++ b/tests/test_service_account.py
@@ -369,6 +369,7 @@ class ServiceAccountCredentialsTests(unittest.TestCase):
 
         self.assertEqual(credentials.access_token, token2)
 
+
 TOKEN_LIFE = service_account._JWTAccessCredentials._MAX_TOKEN_LIFETIME_SECS
 T1 = 42
 T1_DATE = datetime.datetime(1970, 1, 1, second=T1)