Optionally use pickle5 (Redux)

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

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
 =====

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

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)

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

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

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_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') +
  "'))()")

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"