Import pickle from compat in tests

cloudpipe · May 20, 2020 · f356dbd · f356dbd
1 parent d480598
commit f356dbd
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 4 deletions.
diff --git a/tests/cloudpickle_file_test.py b/tests/cloudpickle_file_test.py
@@ -1,7 +1,6 @@
 from __future__ import unicode_literals
 
 import os
-import pickle
 import shutil
 import sys
 import tempfile
@@ -10,6 +9,7 @@
 import pytest
 
 import cloudpickle
+from cloudpickle.compat import pickle
 
 
 class CloudPickleFileTests(unittest.TestCase):

diff --git a/tests/cloudpickle_test.py b/tests/cloudpickle_test.py
@@ -9,7 +9,6 @@
 import logging
 import math
 from operator import itemgetter, attrgetter
-import pickle
 import platform
 import random
 import shutil
@@ -43,6 +42,7 @@
     tornado = None
 
 import cloudpickle
+from cloudpickle.compat import pickle
 from cloudpickle.cloudpickle import _is_dynamic
 from cloudpickle.cloudpickle import _make_empty_cell, cell_set
 from cloudpickle.cloudpickle import _extract_class_dict, _whichmodule
@@ -1228,8 +1228,12 @@ def f():
         # some setup is required to allow pytest apimodules to be correctly
         # serializable.
         from cloudpickle import CloudPickler
-        from cloudpickle import cloudpickle_fast as cp_fast
-        CloudPickler.dispatch_table[type(py.builtin)] = cp_fast._module_reduce
+        if sys.version_info[:2] >= (3, 8):
+            from cloudpickle import cloudpickle_fast as cp_fast
+            CloudPickler.dispatch[
+                type(py.builtin)] = cp_fast._module_reduce
+        else:
+            CloudPickler.dispatch[type(py.builtin)] = CloudPickler.save_module
 
         g = cloudpickle.loads(cloudpickle.dumps(f, protocol=self.protocol))