Add pickling of dict_keys, dict_values, dict_items (#384)

Author: brl0

Diff for: 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 for: 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 for: 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),