Skip to content

Commit

Permalink
Fix pickle for Dataset objects
Browse files Browse the repository at this point in the history
Fixes pydata#167

Related: pydata#166 (though I'm not sure why that didn't work)
  • Loading branch information
shoyer committed Jun 22, 2014
1 parent 75f5f93 commit 306c6d8
Show file tree
Hide file tree
Showing 2 changed files with 16 additions and 15 deletions.
2 changes: 2 additions & 0 deletions test/test_dataset.py
Original file line number Diff line number Diff line change
Expand Up @@ -651,6 +651,8 @@ def test_pickle(self):
data = create_test_data()
roundtripped = pickle.loads(pickle.dumps(data))
self.assertDatasetIdentical(data, roundtripped)
# regression test for #167:
self.assertEqual(data.dimensions, roundtripped.dimensions)

def test_lazy_load(self):
store = InaccessibleVariableDataStore()
Expand Down
29 changes: 14 additions & 15 deletions xray/utils.py
Original file line number Diff line number Diff line change
Expand Up @@ -215,7 +215,18 @@ def ordered_dict_intersection(first_dict, second_dict, compat=equivalent):
return new_dict


class Frozen(Mapping):
class SingleSlotPickleMixin(object):
"""Mixin class to add the ability to pickle objects whose state is defined
by a single __slots__ attribute. Only necessary under Python 2.
"""
def __getstate__(self):
return getattr(self, self.__slots__[0])

def __setstate__(self, state):
setattr(self, self.__slots__[0], state)


class Frozen(Mapping, SingleSlotPickleMixin):
"""Wrapper around an object implementing the mapping interface to make it
immutable. If you really want to modify the mapping, the mutable version is
saved under the `mapping` attribute.
Expand All @@ -240,16 +251,12 @@ def __contains__(self, key):
def __repr__(self):
return '%s(%r)' % (type(self).__name__, self.mapping)

if not PY3:
def __getstate__(self):
return self.__dict__


def FrozenOrderedDict(*args, **kwargs):
return Frozen(OrderedDict(*args, **kwargs))


class SortedKeysDict(MutableMapping):
class SortedKeysDict(MutableMapping, SingleSlotPickleMixin):
"""An wrapper for dictionary-like objects that always iterates over its
items in sorted order by key but is otherwise equivalent to the underlying
mapping.
Expand Down Expand Up @@ -283,12 +290,8 @@ def __repr__(self):
def copy(self):
return type(self)(self.mapping.copy())

if not PY3:
def __getstate__(self):
return self.__dict__


class ChainMap(MutableMapping):
class ChainMap(MutableMapping, SingleSlotPickleMixin):
"""Partial backport of collections.ChainMap from Python>=3.3
Don't return this from any public APIs, since some of the public methods
Expand Down Expand Up @@ -319,10 +322,6 @@ def __iter__(self):
def __len__(self):
raise NotImplementedError

if not PY3:
def __getstate__(self):
return self.__dict__


class NDArrayMixin(object):
"""Mixin class for making wrappers of N-dimensional arrays that conform to
Expand Down

0 comments on commit 306c6d8

Please sign in to comment.