WIP: Optionally use pickle5

cloudpickle/__init__.py

@@ -1,11 +1,10 @@
 from __future__ import absolute_import
 
-import sys
-import pickle
+from cloudpickle.compat import pickle
 
 
 from cloudpickle.cloudpickle import *
-if sys.version_info[:2] >= (3, 8):
+if pickle.HIGHEST_PROTOCOL >= 5:
     from cloudpickle.cloudpickle_fast import CloudPickler, dumps, dump
 
 __version__ = '1.5.0dev0'

cloudpickle/cloudpickle.py

@@ -50,7 +50,7 @@
 import logging
 import opcode
 import operator
-import pickle
+from cloudpickle.compat import pickle
 import platform
 import struct
 import sys
@@ -62,7 +62,6 @@
 from enum import Enum
 
 from typing import Generic, Union, Tuple, Callable
-from pickle import _Pickler as Pickler
 from pickle import _getattribute
 from io import BytesIO
 from importlib._bootstrap import _find_spec
@@ -79,6 +78,9 @@
     ClassVar = None
 
 
+# Aliases using the compat module
+Pickler = pickle._Pickler
+
 # cloudpickle is meant for inter process communication: we expect all
 # communicating processes to run the same Python version hence we favor
 # communication speed over compatibility:

cloudpickle/cloudpickle_fast.py

@@ -15,22 +15,21 @@
 import io
 import itertools
 import logging
-import _pickle
-import pickle
+from cloudpickle.compat import pickle
+from cloudpickle.compat import _pickle
 import sys
 import types
 import weakref
 import typing
 
-from _pickle import Pickler
-
 from .cloudpickle import (
     _is_dynamic, _extract_code_globals, _BUILTIN_TYPE_NAMES, DEFAULT_PROTOCOL,
     _find_imported_submodules, _get_cell_contents, _is_importable_by_name, _builtin_type,
     Enum, _get_or_create_tracker_id,  _make_skeleton_class, _make_skeleton_enum,
     _extract_class_dict, dynamic_subimport, subimport, _typevar_reduce, _get_bases,
 )
 
+Pickler = _pickle.Pickler
 load, loads = _pickle.load, _pickle.loads
 
 
@@ -190,14 +189,23 @@ def _enum_getstate(obj):
 
 def _code_reduce(obj):
     """codeobject reducer"""
-    args = (
-        obj.co_argcount, obj.co_posonlyargcount,
-        obj.co_kwonlyargcount, obj.co_nlocals, obj.co_stacksize,
-        obj.co_flags, obj.co_code, obj.co_consts, obj.co_names,
-        obj.co_varnames, obj.co_filename, obj.co_name,
-        obj.co_firstlineno, obj.co_lnotab, obj.co_freevars,
-        obj.co_cellvars
-    )
+    if hasattr(obj, "co_posonlyargcount"):
+        args = (
+            obj.co_argcount, obj.co_posonlyargcount,
+            obj.co_kwonlyargcount, obj.co_nlocals, obj.co_stacksize,
+            obj.co_flags, obj.co_code, obj.co_consts, obj.co_names,
+            obj.co_varnames, obj.co_filename, obj.co_name,
+            obj.co_firstlineno, obj.co_lnotab, obj.co_freevars,
+            obj.co_cellvars
+        )
+    else:
+        args = (
+            obj.co_argcount, obj.co_kwonlyargcount, obj.co_nlocals,
+            obj.co_stacksize, obj.co_flags, obj.co_code, obj.co_consts,
+            obj.co_names, obj.co_varnames, obj.co_filename,
+            obj.co_name, obj.co_firstlineno, obj.co_lnotab,
+            obj.co_freevars, obj.co_cellvars
+        )
     return types.CodeType, args
 
 
@@ -416,7 +424,8 @@ class CloudPickler(Pickler):
     dispatch[memoryview] = _memoryview_reduce
     dispatch[property] = _property_reduce
     dispatch[staticmethod] = _classmethod_reduce
-    dispatch[types.CellType] = _cell_reduce
+    if hasattr(types, "CellType"):
+        dispatch[types.CellType] = _cell_reduce
     dispatch[types.CodeType] = _code_reduce
     dispatch[types.GetSetDescriptorType] = _getset_descriptor_reduce
     dispatch[types.ModuleType] = _module_reduce
@@ -537,9 +546,11 @@ def _function_getnewargs(self, func):
         # avoid infinite recursion.
         if func.__closure__ is None:
             closure = None
-        else:
+        elif hasattr(types, "CellType"):
             closure = tuple(
                 types.CellType() for _ in range(len(code.co_freevars)))
+        else:
+            closure = list(map(_get_cell_contents, func.__closure__))
 
         return code, base_globals, None, None, closure
 

cloudpickle/compat.py

@@ -0,0 +1,13 @@
+import sys
+
+
+if sys.version_info.major == 3 and sys.version_info.minor < 8:
+    try:
+        import pickle5 as pickle  # noqa: F401
+        import pickle5._pickle as _pickle  # noqa: F401
+    except ImportError:
+        import pickle  # noqa: F401
+        import _pickle  # noqa: F401
+else:
+    import pickle  # noqa: F401
+    import _pickle  # noqa: F401

dev-requirements.txt

@@ -3,6 +3,8 @@ flake8
 pytest
 pytest-cov
 psutil
+# To test on older Python versions
+pickle5; python_version >= '3.6' and python_version <= '3.7'
 # To be able to test tornado coroutines
 tornado
 # To be able to test numpy specific things

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):

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