Merge pull request #168 from shoyer/fix-pickle

Fix pickle for Dataset objects

Author: shoyer

test/test_dataset.py

@@ -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()

xray/utils.py

@@ -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.
@@ -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.
@@ -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
@@ -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