fixes issue python-quantities#113

burnpanck · burnpanck · commit 1b25e4d27adc · 2016-09-06T14:55:46.000+02:00
diff --git a/quantities/quantity.py b/quantities/quantity.py
@@ -697,34 +697,21 @@ def cumprod(self, axis=None, dtype=None, out=None):
 
     # list of unsupported functions: [choose]
 
-    def __getstate__(self):
-        """
-        Return the internal state of the quantity, for pickling
-        purposes.
-
-        """
-        cf = 'CF'[self.flags.fnc]
-        state = (1,
-                 self.shape,
-                 self.dtype,
-                 self.flags.fnc,
-                 self.tostring(cf),
-                 self._dimensionality,
-                 )
-        return state
-
     def __setstate__(self, state):
-        (ver, shp, typ, isf, raw, units) = state
-        np.ndarray.__setstate__(self, (shp, typ, isf, raw))
+        ndarray_state = state[:-1]
+        units = state[-1]
+        np.ndarray.__setstate__(self, ndarray_state)
         self._dimensionality = units
 
     def __reduce__(self):
         """
         Return a tuple for pickling a Quantity.
         """
+        reconstruct,reconstruct_args,state = super(Quantity,self).__reduce__()
+        state = state + (self._dimensionality,)
         return (_reconstruct_quantity,
                 (self.__class__, np.ndarray, (0, ), 'b', ),
-                self.__getstate__())
+                state)
 
     def __deepcopy__(self, memo_dict):
         # constructor copies by default
diff --git a/quantities/tests/test_persistence.py b/quantities/tests/test_persistence.py
@@ -1,6 +1,7 @@
 # -*- coding: utf-8 -*-
 
 import pickle
+import copy
 
 from .. import units as pq
 from ..quantity import Quantity
@@ -47,3 +48,14 @@ def test_uncertainquantity_object_dtype(self):
         y = pickle.loads(pickle.dumps(x))
         self.assertQuantityEqual(x, y)
 
+    def test_copy_quantity(self):
+        for dtype in [float,object]:
+            x = (20*pq.m).astype(dtype)
+            y = pickle.loads(pickle.dumps(x))
+            self.assertQuantityEqual(x, y)
+
+    def test_copy_uncertainquantity(self):
+        for dtype in [float, object]:
+            x = UncertainQuantity(20, 'm', 0.2).astype(dtype)
+            y = pickle.loads(pickle.dumps(x))
+            self.assertQuantityEqual(x, y)
diff --git a/quantities/uncertainquantity.py b/quantities/uncertainquantity.py
@@ -273,18 +273,18 @@ def nanargmin(self,axis=None, out=None):
     def nanargmax(self,axis=None, out=None):
         return np.nanargmax(self.magnitude)
 
-    def __getstate__(self):
-        """
-        Return the internal state of the quantity, for pickling
-        purposes.
-
-        """
-        state = list(super(UncertainQuantity, self).__getstate__())
-        state.append(self._uncertainty)
-        return tuple(state)
-
     def __setstate__(self, state):
-        (ver, shp, typ, isf, raw, units, sigma) = state
-        np.ndarray.__setstate__(self, (shp, typ, isf, raw))
+        ndarray_state = state[:-2]
+        units, sigma = state[-2:]
+        np.ndarray.__setstate__(self, ndarray_state)
         self._dimensionality = units
         self._uncertainty = sigma
+
+    def __reduce__(self):
+        """
+        Return a tuple for pickling a Quantity.
+        """
+        reconstruct, reconstruct_args, state = super(UncertainQuantity, self).__reduce__()
+        state = state + (self._uncertainty,)
+        return reconstruct, reconstruct_args, state
+
