From c27d139bbe6315a9bc46e9730001dd2d2e13e420 Mon Sep 17 00:00:00 2001 From: Siyuan Date: Wed, 11 Sep 2019 18:05:47 -0700 Subject: [PATCH 1/3] add the missing buffer_callback argument --- cloudpickle/cloudpickle_fast.py | 12 ++++++------ 1 file changed, 6 insertions(+), 6 deletions(-) diff --git a/cloudpickle/cloudpickle_fast.py b/cloudpickle/cloudpickle_fast.py index dfe554d62..691bc77cc 100644 --- a/cloudpickle/cloudpickle_fast.py +++ b/cloudpickle/cloudpickle_fast.py @@ -34,7 +34,7 @@ # Shorthands similar to pickle.dump/pickle.dumps -def dump(obj, file, protocol=None): +def dump(obj, file, protocol=None, buffer_callback=None): """Serialize obj as bytes streamed into file protocol defaults to cloudpickle.DEFAULT_PROTOCOL which is an alias to @@ -44,10 +44,10 @@ def dump(obj, file, protocol=None): Set protocol=pickle.DEFAULT_PROTOCOL instead if you need to ensure compatibility with older versions of Python. """ - CloudPickler(file, protocol=protocol).dump(obj) + CloudPickler(file, protocol=protocol, buffer_callback=buffer_callback).dump(obj) -def dumps(obj, protocol=None): +def dumps(obj, protocol=None, buffer_callback=None): """Serialize obj as a string of bytes allocated in memory protocol defaults to cloudpickle.DEFAULT_PROTOCOL which is an alias to @@ -58,7 +58,7 @@ def dumps(obj, protocol=None): compatibility with older versions of Python. """ with io.BytesIO() as file: - cp = CloudPickler(file, protocol=protocol) + cp = CloudPickler(file, protocol=protocol, buffer_callback=buffer_callback) cp.dump(obj) return file.getvalue() @@ -415,10 +415,10 @@ class CloudPickler(Pickler): dispatch[types.MappingProxyType] = _mappingproxy_reduce dispatch[weakref.WeakSet] = _weakset_reduce - def __init__(self, file, protocol=None): + def __init__(self, file, protocol=None, buffer_callback=None): if protocol is None: protocol = DEFAULT_PROTOCOL - Pickler.__init__(self, file, protocol=protocol) + Pickler.__init__(self, file, protocol=protocol, buffer_callback=buffer_callback) # map functions __globals__ attribute ids, to ensure that functions # sharing the same global namespace at pickling time also share their # global namespace at unpickling time. From 7a1de7044cd2e91533871661bf2ebfbacb2d1016 Mon Sep 17 00:00:00 2001 From: Olivier Grisel Date: Fri, 7 Feb 2020 15:20:29 +0100 Subject: [PATCH 2/3] Add test for out-of-band pickling --- tests/cloudpickle_test.py | 16 ++++++++++++++++ 1 file changed, 16 insertions(+) diff --git a/tests/cloudpickle_test.py b/tests/cloudpickle_test.py index 8bd210abe..3b1378e80 100644 --- a/tests/cloudpickle_test.py +++ b/tests/cloudpickle_test.py @@ -1876,6 +1876,22 @@ def __getattr__(self, name): with pytest.raises(pickle.PicklingError, match='recursion'): cloudpickle.dumps(a) + def test_out_of_band_buffers(self): + if self.protocol < 5: + pytest.skip("Need Pickle Protocol 5 or later") + np = pytest.importorskip("numpy") + + class LocallyDefinedClass: + data = np.zeros(10) + + data_instance = LocallyDefinedClass() + buffers = [] + pickle_bytes = cloudpickle.dumps(data_instance, protocol=self.protocol, + buffer_callback=buffers.append) + assert len(buffers) == 1 + reconstructed = pickle.loads(pickle_bytes, buffers=buffers) + np.testing.assert_allclose(reconstructed.data, data_instance.data) + class Protocol2CloudPickleTest(CloudPickleTest): From 79ec004388baa14d602f1acff66680ef55fbea69 Mon Sep 17 00:00:00 2001 From: Olivier Grisel Date: Fri, 7 Feb 2020 15:23:28 +0100 Subject: [PATCH 3/3] Add changelog entry --- CHANGES.md | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/CHANGES.md b/CHANGES.md index 94dcc014f..f2ae5bcd7 100644 --- a/CHANGES.md +++ b/CHANGES.md @@ -1,6 +1,10 @@ 1.2.3 ===== +- Add support for out-of-band pickling (Python 3.8 and later). + https://docs.python.org/3/library/pickle.html#example + ([issue #308](https://github.com/cloudpipe/cloudpickle/pull/308)) + 1.2.2 =====