Optionally use pickle5 (Redux) (#370)

Co-authored-by: Pierre Glaser <[email protected]>

Author: jakirkham

Diff for: .github/workflows/testing.yml

@@ -228,6 +228,7 @@ jobs:
         git clone https://github.com/ray-project/ray.git ../ray
         cp -R ../ray/python/ray/tests $PROJECT_DIR/tests
         cp cloudpickle/cloudpickle.py $PROJECT_DIR/cloudpickle/cloudpickle.py
+        cp cloudpickle/compat.py $PROJECT_DIR/cloudpickle/compat.py
         cp cloudpickle/cloudpickle_fast.py $PROJECT_DIR/cloudpickle/cloudpickle_fast.py
     - name: Test the downstream project
       run: |

Diff for: CHANGES.md

@@ -5,6 +5,10 @@
   importable modules.
   ([issue #360](https://github.com/cloudpipe/cloudpickle/issues/354))
 
+- Add optional dependency on `pickle5` to get improved performance on
+  Python 3.6 and 3.7.
+  ([PR #370](https://github.com/cloudpipe/cloudpickle/pull/370))
+
 
 1.4.1
 =====

Diff for: cloudpickle/cloudpickle.py

@@ -45,7 +45,6 @@
 import builtins
 import dis
 import opcode
-import pickle
 import platform
 import sys
 import types
@@ -55,6 +54,7 @@
 import typing
 import warnings
 
+from .compat import pickle
 from typing import Generic, Union, Tuple, Callable
 from pickle import _getattribute
 from importlib._bootstrap import _find_spec

Diff for: cloudpickle/cloudpickle_fast.py

@@ -15,7 +15,6 @@
 import io
 import itertools
 import logging
-import pickle
 import sys
 import struct
 import types
@@ -25,6 +24,7 @@
 from enum import Enum
 from collections import ChainMap
 
+from .compat import pickle, Pickler
 from .cloudpickle import (
     _extract_code_globals, _BUILTIN_TYPE_NAMES, DEFAULT_PROTOCOL,
     _find_imported_submodules, _get_cell_contents, _is_importable,
@@ -37,8 +37,8 @@
 
 )
 
-if sys.version_info >= (3, 8) and not PYPY:
-    from _pickle import Pickler
+
+if pickle.HIGHEST_PROTOCOL >= 5 and not PYPY:
     # Shorthands similar to pickle.dump/pickle.dumps
 
     def dump(obj, file, protocol=None, buffer_callback=None):
@@ -73,8 +73,6 @@ def dumps(obj, protocol=None, buffer_callback=None):
             return file.getvalue()
 
 else:
-    from pickle import _Pickler as Pickler
-
     # Shorthands similar to pickle.dump/pickle.dumps
     def dump(obj, file, protocol=None):
         """Serialize obj as bytes streamed into file
@@ -551,6 +549,17 @@ def dump(self, obj):
                 raise
 
     if pickle.HIGHEST_PROTOCOL >= 5:
+        # `CloudPickler.dispatch` is only left for backward compatibility - note
+        # that when using protocol 5, `CloudPickler.dispatch` is not an
+        # extension of `Pickler.dispatch` dictionary, because CloudPickler
+        # subclasses the C-implemented Pickler, which does not expose a
+        # `dispatch` attribute.  Earlier versions of the protocol 5 CloudPickler
+        # used `CloudPickler.dispatch` as a class-level attribute storing all
+        # reducers implemented by cloudpickle, but the attribute name was not a
+        # great choice given the meaning of `Cloudpickler.dispatch` when
+        # `CloudPickler` extends the pure-python pickler.
+        dispatch = dispatch_table
+
         # Implementation of the reducer_override callback, in order to
         # efficiently serialize dynamic functions and classes by subclassing
         # the C-implemented Pickler.
@@ -604,6 +613,11 @@ def reducer_override(self, obj):
               reducers, such as Exceptions. See
               https://github.com/cloudpipe/cloudpickle/issues/248
             """
+            if sys.version_info[:2] < (3, 7) and _is_parametrized_type_hint(obj):  # noqa  # pragma: no branch
+                return (
+                    _create_parametrized_type_hint,
+                    parametrized_type_hint_getinitargs(obj)
+                )
             t = type(obj)
             try:
                 is_anyclass = issubclass(t, type)

Diff for: cloudpickle/compat.py

@@ -0,0 +1,13 @@
+import sys
+
+
+if sys.version_info < (3, 8):
+    try:
+        import pickle5 as pickle  # noqa: F401
+        from pickle5 import Pickler  # noqa: F401
+    except ImportError:
+        import pickle  # noqa: F401
+        from pickle import _Pickler as Pickler  # noqa: F401
+else:
+    import pickle  # noqa: F401
+    from _pickle import Pickler  # noqa: F401

Diff for: dev-requirements.txt

@@ -3,11 +3,13 @@ flake8
 pytest
 pytest-cov
 psutil
+# To test on older Python versions
+pickle5 >=0.0.11 ; python_version <= '3.7' and python_implementation == 'CPython'
 # To be able to test tornado coroutines
 tornado
 # To be able to test numpy specific things
 # but do not build numpy from source on Python nightly
-numpy; python_version <= '3.8'
+numpy >=1.18.5; python_version <= '3.8'
 # Code coverage uploader for Travis:
 codecov
 coverage

Diff for: 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 for: 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_importable
 from cloudpickle.cloudpickle import _make_empty_cell, cell_set
 from cloudpickle.cloudpickle import _extract_class_dict, _whichmodule
@@ -521,7 +521,7 @@ def test_module_locals_behavior(self):
         pickled_func_path = os.path.join(self.tmpdir, 'local_func_g.pkl')
 
         child_process_script = '''
-        import pickle
+        from cloudpickle.compat import pickle
         import gc
         with open("{pickled_func_path}", 'rb') as f:
             func = pickle.load(f)
@@ -606,7 +606,7 @@ def test_load_dynamic_module_in_grandchild_process(self):
         child_process_module_file = os.path.join(
             self.tmpdir, 'dynamic_module_from_child_process.pkl')
         child_process_script = '''
-            import pickle
+            from cloudpickle.compat import pickle
             import textwrap
 
             import cloudpickle
@@ -626,7 +626,7 @@ def test_load_dynamic_module_in_grandchild_process(self):
 
         # The script ran by the process created by the child process
         child_of_child_process_script = """ '''
-                import pickle
+                from cloudpickle.compat import pickle
                 with open('{child_process_module_file}','rb') as fid:
                     mod = pickle.load(fid)
                 ''' """
@@ -681,7 +681,7 @@ def my_small_function(x, y):
         assert b'math' not in b
 
     def test_module_importability(self):
-        import pickle  # decouple this test from global imports
+        from cloudpickle.compat import pickle
         import os.path
         import distutils
         import distutils.ccompiler
@@ -1008,7 +1008,8 @@ def example():
 
         # choose "subprocess" rather than "multiprocessing" because the latter
         # library uses fork to preserve the parent environment.
-        command = ("import pickle, base64; "
+        command = ("import base64; "
+                   "from cloudpickle.compat import pickle; "
                    "pickle.loads(base64.b32decode('" +
                    base64.b32encode(s).decode('ascii') +
                    "'))()")
@@ -1029,7 +1030,8 @@ def example():
 
         s = cloudpickle.dumps(example, protocol=self.protocol)
 
-        command = ("import pickle, base64; "
+        command = ("import base64; "
+                   "from cloudpickle.compat import pickle; "
                    "pickle.loads(base64.b32decode('" +
                    base64.b32encode(s).decode('ascii') +
                    "'))()")

Diff for: tests/testutils.py

@@ -4,14 +4,15 @@
 import tempfile
 import base64
 from subprocess import Popen, check_output, PIPE, STDOUT, CalledProcessError
-from pickle import loads
+from cloudpickle.compat import pickle
 from contextlib import contextmanager
 from concurrent.futures import ProcessPoolExecutor
 
 import psutil
 from cloudpickle import dumps
 from subprocess import TimeoutExpired
 
+loads = pickle.loads
 TIMEOUT = 60
 TEST_GLOBALS = "a test value"