diff --git a/cloudpickle/cloudpickle.py b/cloudpickle/cloudpickle.py index 08882306..66bb3f94 100644 --- a/cloudpickle/cloudpickle.py +++ b/cloudpickle/cloudpickle.py @@ -96,6 +96,10 @@ PYPY = platform.python_implementation() == "PyPy" +# Keep using the pure-Python pickler on PyPy: its accelerated reducer_override +# does not support the state setters used by cloudpickle. +_PicklerBase = pickle._Pickler if PYPY else pickle.Pickler + builtin_code_type = None if PYPY: # builtin-code objects only exist in pypy @@ -1235,7 +1239,7 @@ def _get_dataclass_field_type_sentinel(name): return _DATACLASSE_FIELD_TYPE_SENTINELS[name] -class Pickler(pickle.Pickler): +class Pickler(_PicklerBase): # set of reducers defined and used by cloudpickle (private) _dispatch_table = {} _dispatch_table[classmethod] = _classmethod_reduce @@ -1339,16 +1343,15 @@ def __init__(self, file, protocol=None, buffer_callback=None): self.globals_ref = {} self.proto = int(protocol) - if not PYPY: - # pickle.Pickler is the C implementation of the CPython pickler and - # therefore we rely on reduce_override method to customize the pickler - # behavior. + if not hasattr(_PicklerBase, "dispatch"): + # Accelerated Pickler implementations do not expose a dispatch + # dictionary. Rely on reducer_override to customize their behavior. # `cloudpickle.Pickler.dispatch` is only left for backward # compatibility - note that when using protocol 5, # `cloudpickle.Pickler.dispatch` is not an extension of # `pickle._Pickler.dispatch` dictionary, because `cloudpickle.Pickler` - # subclasses the C-implemented `pickle.Pickler`, which does not expose + # subclasses the accelerated `pickle.Pickler`, which does not expose # a `dispatch` attribute. Earlier versions of `cloudpickle.Pickler` # used `cloudpickle.Pickler.dispatch` as a class-level attribute # storing all reducers implemented by cloudpickle, but the attribute @@ -1359,12 +1362,7 @@ def __init__(self, file, protocol=None, buffer_callback=None): # Implementation of the reducer_override callback, in order to # efficiently serialize dynamic functions and classes by subclassing - # the C-implemented `pickle.Pickler`. - # TODO: decorrelate reducer_override (which is tied to CPython's - # implementation - would it make sense to backport it to pypy? - and - # pickle's protocol 5 which is implementation agnostic. Currently, the - # availability of both notions coincide on CPython's pickle, but it may - # not be the case anymore when pypy implements protocol 5. + # the accelerated `pickle.Pickler`. def reducer_override(self, obj): """Type-agnostic reducing callback for function and classes. @@ -1413,11 +1411,10 @@ def reducer_override(self, obj): return NotImplemented else: - # When reducer_override is not available, hack the pure-Python - # Pickler's types.FunctionType and type savers. Note: the type saver - # must override Pickler.save_global, because pickle.py contains a - # hard-coded call to save_global when pickling meta-classes. - dispatch = pickle.Pickler.dispatch.copy() + # For pure-Python Picklers, customize the types.FunctionType and type + # savers. The type saver must override Pickler.save_global, because + # pickle.py contains a hard-coded call to save_global for meta-classes. + dispatch = _PicklerBase.dispatch.copy() def _save_reduce_pickle5( self, diff --git a/tests/cloudpickle_test.py b/tests/cloudpickle_test.py index e2097d1c..66581a03 100644 --- a/tests/cloudpickle_test.py +++ b/tests/cloudpickle_test.py @@ -87,6 +87,50 @@ def _maybe_remove(list_, item): return list_ +@pytest.mark.parametrize("pickler_name", ["Pickler", "_Pickler"]) +@pytest.mark.parametrize("protocol", range(pickle.HIGHEST_PROTOCOL + 1)) +def test_pickle_implementations(pickler_name, protocol): + # Select the implementation before importing cloudpickle, as PyPy switched + # from the pure-Python Pickler to an accelerated one in PyPy 7.3.22. + code = f""" +import collections +import io +import pickle +pickle.Pickler = getattr(pickle, {pickler_name!r}) +import cloudpickle + +class CustomMeta(type): + pass + +class DynamicClass(metaclass=CustomMeta): + value = 42 + +objects = [pickle.whichmodule, collections.Counter, collections.Counter("ababa"), + {{"value": 42}}, b"\\x00\\xff", DynamicClass, DynamicClass()] +restored = pickle.loads(cloudpickle.dumps(objects, protocol={protocol})) +assert restored[0] is pickle.whichmodule +assert restored[1] is collections.Counter +assert restored[:5] == objects[:5] +assert isinstance(restored[5], CustomMeta) +assert isinstance(restored[6], restored[5]) +assert restored[6].value == 42 + +if hasattr(cloudpickle.Pickler, "reducer_override"): + sentinel = lambda: None + + class CustomPickler(cloudpickle.Pickler): + def _function_reduce(self, obj): + if obj is sentinel: + return str, ("hook",) + return super()._function_reduce(obj) + + stream = io.BytesIO() + CustomPickler(stream, protocol={protocol}).dump(sentinel) + assert pickle.loads(stream.getvalue()) == "hook" +""" + assert_run_python_script(code) + + def test_extract_class_dict(): class A(int): """A docstring"""