Skip to content

Commit 0c62ae0

Browse files
ogriselsmacke
andauthored
allow pickling / depickling of odict_keys, odict_values, odict_items (#429)
Co-authored-by: Stephen Macke <[email protected]>
1 parent 685bf0c commit 0c62ae0

File tree

4 files changed

+60
-7
lines changed

4 files changed

+60
-7
lines changed

Diff for: CHANGES.md

+4
Original file line numberDiff line numberDiff line change
@@ -17,6 +17,10 @@ dev
1717
_is_parametrized_type_hint to limit false positives.
1818
([PR #409](https://github.com/cloudpipe/cloudpickle/pull/409))
1919

20+
- Support pickling / depickling of OrderedDict KeysView, ValuesView, and
21+
ItemsView, following similar strategy for vanilla Python dictionaries.
22+
([PR #423](https://github.com/cloudpipe/cloudpickle/pull/423))
23+
2024
- Suppressed a source of non-determinism when pickling dynamically defined
2125
functions and handles the deprecation of co_lnotab in Python 3.10+.
2226
([PR #428](https://github.com/cloudpipe/cloudpickle/pull/428))

Diff for: cloudpickle/cloudpickle.py

+16-6
Original file line numberDiff line numberDiff line change
@@ -55,6 +55,7 @@
5555
import warnings
5656

5757
from .compat import pickle
58+
from collections import OrderedDict
5859
from typing import Generic, Union, Tuple, Callable
5960
from pickle import _getattribute
6061
from importlib._bootstrap import _find_spec
@@ -855,13 +856,22 @@ def _get_bases(typ):
855856
return getattr(typ, bases_attr)
856857

857858

858-
def _make_dict_keys(obj):
859-
return dict.fromkeys(obj).keys()
859+
def _make_dict_keys(obj, is_ordered=False):
860+
if is_ordered:
861+
return OrderedDict.fromkeys(obj).keys()
862+
else:
863+
return dict.fromkeys(obj).keys()
860864

861865

862-
def _make_dict_values(obj):
863-
return {i: _ for i, _ in enumerate(obj)}.values()
866+
def _make_dict_values(obj, is_ordered=False):
867+
if is_ordered:
868+
return OrderedDict((i, _) for i, _ in enumerate(obj)).values()
869+
else:
870+
return {i: _ for i, _ in enumerate(obj)}.values()
864871

865872

866-
def _make_dict_items(obj):
867-
return obj.items()
873+
def _make_dict_items(obj, is_ordered=False):
874+
if is_ordered:
875+
return OrderedDict(obj).items()
876+
else:
877+
return obj.items()

Diff for: cloudpickle/cloudpickle_fast.py

+22-1
Original file line numberDiff line numberDiff line change
@@ -23,7 +23,7 @@
2323
import typing
2424

2525
from enum import Enum
26-
from collections import ChainMap
26+
from collections import ChainMap, OrderedDict
2727

2828
from .compat import pickle, Pickler
2929
from .cloudpickle import (
@@ -437,6 +437,24 @@ def _dict_items_reduce(obj):
437437
return _make_dict_items, (dict(obj), )
438438

439439

440+
def _odict_keys_reduce(obj):
441+
# Safer not to ship the full dict as sending the rest might
442+
# be unintended and could potentially cause leaking of
443+
# sensitive information
444+
return _make_dict_keys, (list(obj), True)
445+
446+
447+
def _odict_values_reduce(obj):
448+
# Safer not to ship the full dict as sending the rest might
449+
# be unintended and could potentially cause leaking of
450+
# sensitive information
451+
return _make_dict_values, (list(obj), True)
452+
453+
454+
def _odict_items_reduce(obj):
455+
return _make_dict_items, (dict(obj), True)
456+
457+
440458
# COLLECTIONS OF OBJECTS STATE SETTERS
441459
# ------------------------------------
442460
# state setters are called at unpickling time, once the object is created and
@@ -513,6 +531,9 @@ class CloudPickler(Pickler):
513531
_dispatch_table[_collections_abc.dict_keys] = _dict_keys_reduce
514532
_dispatch_table[_collections_abc.dict_values] = _dict_values_reduce
515533
_dispatch_table[_collections_abc.dict_items] = _dict_items_reduce
534+
_dispatch_table[type(OrderedDict().keys())] = _odict_keys_reduce
535+
_dispatch_table[type(OrderedDict().values())] = _odict_values_reduce
536+
_dispatch_table[type(OrderedDict().items())] = _odict_items_reduce
516537

517538

518539
dispatch_table = ChainMap(_dispatch_table, copyreg.dispatch_table)

Diff for: tests/cloudpickle_test.py

+18
Original file line numberDiff line numberDiff line change
@@ -229,6 +229,24 @@ def test_dict_items(self):
229229
self.assertEqual(results, items)
230230
assert isinstance(results, _collections_abc.dict_items)
231231

232+
def test_odict_keys(self):
233+
keys = collections.OrderedDict([("a", 1), ("b", 2)]).keys()
234+
results = pickle_depickle(keys)
235+
self.assertEqual(results, keys)
236+
assert type(keys) == type(results)
237+
238+
def test_odict_values(self):
239+
values = collections.OrderedDict([("a", 1), ("b", 2)]).values()
240+
results = pickle_depickle(values)
241+
self.assertEqual(list(results), list(values))
242+
assert type(values) == type(results)
243+
244+
def test_odict_items(self):
245+
items = collections.OrderedDict([("a", 1), ("b", 2)]).items()
246+
results = pickle_depickle(items)
247+
self.assertEqual(results, items)
248+
assert type(items) == type(results)
249+
232250
def test_sliced_and_non_contiguous_memoryview(self):
233251
buffer_obj = memoryview(b"Hello!" * 3)[2:15:2]
234252
self.assertEqual(pickle_depickle(buffer_obj, protocol=self.protocol),

0 commit comments

Comments
 (0)