Added explicit pickle test and fixed curve.Data pickling.

[hooke.git] / hooke / curve.py

diff --git a/hooke/curve.py b/hooke/curve.py
index 816a921a86470f842ee2bda1d7d2131a73722de0..146acfa8abf18928f823672ff4e774addd4eaed0 100644 (file)
--- a/hooke/curve.py
+++ b/hooke/curve.py
@@ -26,9 +26,16 @@ import numpy
 
 class NotRecognized (ValueError):
     def __init__(self, curve):
-        msg = 'Not a recognizable curve format: %s' % curve.path
-        ValueError.__init__(self, msg)
-        self.curve = curve
+        self.__setstate__(curve)
+
+    def __getstate__(self):
+        return self.curve
+
+    def __setstate__(self, data):
+        if isinstance(data, Curve):
+            msg = 'Not a recognizable curve format: %s' % data.path
+            super(NotRecognized, self).__init__(msg)
+            self.curve = data
 
 class Data (numpy.ndarray):
     """Stores a single, continuous data set.
@@ -37,16 +44,52 @@ class Data (numpy.ndarray):
 
     See :mod:`numpy.doc.subclassing` for the peculiarities of
     subclassing :class:`numpy.ndarray`.
+
+    Examples
+    --------
+
+    >>> d = Data(shape=(3,2), info={'columns':['distance (m)', 'force (N)']})
+    >>> type(d)
+    <class 'hooke.curve.Data'>
+    >>> for i in range(3): # initialize d
+    ...    for j in range(2):
+    ...        d[i,j] = i*10 + j
+    >>> d
+    Data([[  0.,   1.],
+           [ 10.,  11.],
+           [ 20.,  21.]])
+    >>> d.info
+    {'columns': ['distance (m)', 'force (N)']}
+
+    The information gets passed on to slices.
+
+    >>> row_a = d[:,0]
+    >>> row_a
+    Data([  0.,  10.,  20.])
+    >>> row_a.info
+    {'columns': ['distance (m)', 'force (N)']}
+
+    The data-type is also pickleable, to ensure we can move it between
+    processes with :class:`multiprocessing.Queue`\s.
+
+    >>> import pickle
+    >>> s = pickle.dumps(d)
+    >>> z = pickle.loads(s)
+    >>> z
+    Data([[  0.,   1.],
+           [ 10.,  11.],
+           [ 20.,  21.]])
+    >>> z.info
+    {'columns': ['distance (m)', 'force (N)']}
     """
-    def __new__(self, subtype, shape, dtype=numpy.float, buffer=None, offset=0,
+    def __new__(subtype, shape, dtype=numpy.float, buffer=None, offset=0,
                 strides=None, order=None, info=None):
         """Create the ndarray instance of our type, given the usual
         input arguments.  This will call the standard ndarray
         constructor, but return an object of our type.
         """
-        obj = np.ndarray.__new__(subtype=subtype, shape=shape, dtype=dtype,
-                                 buffer=buffer, offset=offset, strides=strides,
-                                 order=order)
+        obj = numpy.ndarray.__new__(
+            subtype, shape, dtype, buffer, offset, strides, order)
         # add the new attribute to the created instance
         if info == None:
             info = {}
@@ -61,6 +104,34 @@ class Data (numpy.ndarray):
         self.info = getattr(obj, 'info', {})
         # We do not need to return anything
 
+    def __reduce__(self):
+        """Collapse an instance for pickling.
+
+        Returns
+        -------
+        reconstruct : callable
+            Called to create the initial version of the object.
+        args : tuple
+            A tuple of arguments for `reconstruct`
+        state : (optional)
+            The state to be passed to __setstate__, if present.
+        iter : iterator (optional)
+            Yielded items will be appended to the reconstructed
+            object.
+        dict : iterator (optional)
+            Yielded (key,value) tuples pushed back onto the
+            reconstructed object.
+        """
+        base_reduce = list(numpy.ndarray.__reduce__(self))
+        # tack our stuff onto ndarray's setstate portion.
+        base_reduce[2] = (base_reduce[2], (self.info,))
+        return tuple(base_reduce)
+
+    def __setstate__(self, state):
+        base_class_state,own_state = state
+        numpy.ndarray.__setstate__(self, base_class_state)
+        self.info, = own_state
+
 
 class Curve (object):
     """A grouped set of :class:`Data` runs from the same file with metadata.
@@ -87,7 +158,7 @@ class Curve (object):
         #the data dictionary contains: {name of data: list of data sets [{[x], [y]}]
         self.path = path
         self.driver = None
-        self.data = []
+        self.data = None
         if info == None:
             info = {}
         self.info = info
@@ -106,4 +177,12 @@ class Curve (object):
     def load(self):
         """Use the driver to read the curve into memory.
         """
-        pass
+        data,info = self.driver.read(self.path)
+        self.data = data
+        for key,value in info.items():
+            self.info[key] = value
+
+    def unload(self):
+        """Release memory intensive :attr:`.data`.
+        """
+        self.data = None