diff --git a/cloudpickle/cloudpickle.py b/cloudpickle/cloudpickle.py index 8e683e7a6..05d52afa0 100644 --- a/cloudpickle/cloudpickle.py +++ b/cloudpickle/cloudpickle.py @@ -828,3 +828,15 @@ def _get_bases(typ): # For regular class objects bases_attr = '__bases__' return getattr(typ, bases_attr) + + +def _make_dict_keys(obj): + return dict.fromkeys(obj).keys() + + +def _make_dict_values(obj): + return {i: _ for i, _ in enumerate(obj)}.values() + + +def _make_dict_items(obj): + return obj.items() diff --git a/cloudpickle/cloudpickle_fast.py b/cloudpickle/cloudpickle_fast.py index e8e46b88f..fa8da0f63 100644 --- a/cloudpickle/cloudpickle_fast.py +++ b/cloudpickle/cloudpickle_fast.py @@ -10,6 +10,7 @@ guards present in cloudpickle.py that were written to handle PyPy specificities are not present in cloudpickle_fast.py """ +import _collections_abc import abc import copyreg import io @@ -33,8 +34,8 @@ _typevar_reduce, _get_bases, _make_cell, _make_empty_cell, CellType, _is_parametrized_type_hint, PYPY, cell_set, parametrized_type_hint_getinitargs, _create_parametrized_type_hint, - builtin_code_type - + builtin_code_type, + _make_dict_keys, _make_dict_values, _make_dict_items, ) @@ -400,6 +401,24 @@ def _class_reduce(obj): return NotImplemented +def _dict_keys_reduce(obj): + # Safer not to ship the full dict as sending the rest might + # be unintended and could potentially cause leaking of + # sensitive information + return _make_dict_keys, (list(obj), ) + + +def _dict_values_reduce(obj): + # Safer not to ship the full dict as sending the rest might + # be unintended and could potentially cause leaking of + # sensitive information + return _make_dict_values, (list(obj), ) + + +def _dict_items_reduce(obj): + return _make_dict_items, (dict(obj), ) + + # COLLECTIONS OF OBJECTS STATE SETTERS # ------------------------------------ # state setters are called at unpickling time, once the object is created and @@ -473,6 +492,10 @@ class CloudPickler(Pickler): _dispatch_table[types.MappingProxyType] = _mappingproxy_reduce _dispatch_table[weakref.WeakSet] = _weakset_reduce _dispatch_table[typing.TypeVar] = _typevar_reduce + _dispatch_table[_collections_abc.dict_keys] = _dict_keys_reduce + _dispatch_table[_collections_abc.dict_values] = _dict_values_reduce + _dispatch_table[_collections_abc.dict_items] = _dict_items_reduce + dispatch_table = ChainMap(_dispatch_table, copyreg.dispatch_table) diff --git a/tests/cloudpickle_test.py b/tests/cloudpickle_test.py index ff5d03d94..634d70a5a 100644 --- a/tests/cloudpickle_test.py +++ b/tests/cloudpickle_test.py @@ -1,5 +1,6 @@ from __future__ import division +import _collections_abc import abc import collections import base64 @@ -31,7 +32,7 @@ # tests should be skipped if these modules are not available import numpy as np import scipy.special as spp -except ImportError: +except (ImportError, RuntimeError): np = None spp = None @@ -207,6 +208,24 @@ def test_memoryview(self): self.assertEqual(pickle_depickle(buffer_obj, protocol=self.protocol), buffer_obj.tobytes()) + def test_dict_keys(self): + keys = {"a": 1, "b": 2}.keys() + results = pickle_depickle(keys) + self.assertEqual(results, keys) + assert isinstance(results, _collections_abc.dict_keys) + + def test_dict_values(self): + values = {"a": 1, "b": 2}.values() + results = pickle_depickle(values) + self.assertEqual(sorted(results), sorted(values)) + assert isinstance(results, _collections_abc.dict_values) + + def test_dict_items(self): + items = {"a": 1, "b": 2}.items() + results = pickle_depickle(items) + self.assertEqual(results, items) + assert isinstance(results, _collections_abc.dict_items) + def test_sliced_and_non_contiguous_memoryview(self): buffer_obj = memoryview(b"Hello!" * 3)[2:15:2] self.assertEqual(pickle_depickle(buffer_obj, protocol=self.protocol),