Fix python3 cPickle import errors

apache · Feb 4, 2017 · ce50e6e · ce50e6e
1 parent 167ed33
commit ce50e6e
Show file tree

Hide file tree

Showing 2 changed files with 14 additions and 7 deletions.
diff --git a/superset/results_backends.py b/superset/results_backends.py
@@ -7,7 +7,11 @@
 from __future__ import print_function
 from __future__ import unicode_literals
 
-import cPickle
+try:
+   import cPickle as pickle
+except:
+   import pickle
+
 import logging
 import StringIO
 
@@ -63,7 +67,7 @@ def get(self, key):
                 return None
             else:
                 value_file.seek(0)
-                return cPickle.load(value_file)
+                return pickle.load(value_file)
 
     def delete(self, key):
         """Delete `key` from the cache.
@@ -110,7 +114,7 @@ def set(self, key, value, timeout=None):
         :rtype: boolean
         """
         value_file = StringIO.StringIO()
-        cPickle.dump(value, value_file)
+        pickle.dump(value, value_file)
 
         try:
             value_file.seek(0)

diff --git a/tests/results_backends_tests.py b/tests/results_backends_tests.py
@@ -1,4 +1,7 @@
-import cPickle
+try:
+    import cPickle as pickle
+except:
+    import pickle
 
 import mock
 
@@ -24,7 +27,7 @@ def setUp(self, mock_boto3_client):
 
     @staticmethod
     def _mock_download_fileobj(bucket, key, value_file):
-        value_file.write(cPickle.dumps('%s:%s' % (bucket, key)))
+        value_file.write(pickle.dumps('%s:%s' % (bucket, key)))
 
     @staticmethod
     def _mock_key_exists(key):
@@ -41,7 +44,7 @@ def test_s3_cache_set(self):
 
         call_args = self.mock_s3_client.upload_fileobj.call_args_list[0][0]
 
-        self.assertEquals(cPickle.loads(call_args[0].getvalue()), 'test-value')
+        self.assertEquals(pickle.loads(call_args[0].getvalue()), 'test-value')
         self.assertEquals(call_args[1], 'test-bucket')
         self.assertEquals(call_args[2], 'test-prefix/test-key')
 
@@ -109,7 +112,7 @@ def test_s3_cache_add_does_not_exist(self):
 
         call_args = self.mock_s3_client.upload_fileobj.call_args_list[0][0]
 
-        self.assertEquals(cPickle.loads(call_args[0].getvalue()), 'test-value')
+        self.assertEquals(pickle.loads(call_args[0].getvalue()), 'test-value')
         self.assertEquals(call_args[1], 'test-bucket')
         self.assertEquals(call_args[2], 'test-prefix/test-key2')