Add support for {,u}int{16,32} segments in JPK driver.

TODO: check that these are converted to float32 or similar when
converting from bits to meters.  Potentially change from multi-column
Data blocks to lists of single column blocks to allow for different
dtypes in each block.

diff --git a/hooke/driver/jpk.py b/hooke/driver/jpk.py
index fec6cf523b01de9955de74f012c8da79ce90f50e..948d75abc556aab463da004d4879add94ef9f70b 100644 (file)
--- a/hooke/driver/jpk.py
+++ b/hooke/driver/jpk.py
@@ -127,16 +127,42 @@ class JPKDriver (Driver):
         return self._zip_scale_segment(d, path, info)
 
     def _zip_channel(self, zipfile, segment_index, channel_name, chan_info):
-        f = zipfile.open(os.path.join(
-                'segments', str(segment_index),
-                chan_info['data']['file']['name']), 'r')
-        assert chan_info['data']['file']['format'] == 'raw', \
-            'Non-raw data format:\n%s' % pprint.pformat(chan_info)
-        assert chan_info['data']['type'] == 'float-data', \
-            'Non-float data format:\n%s' % pprint.pformat(chan_info)
-        data = numpy.frombuffer(
-            buffer(f.read()),
-            dtype=numpy.dtype(numpy.float32).newbyteorder('>'))
+        with Closing(zipfile.open(os.path.join(
+                    'segments', str(segment_index),
+                    chan_info['data']['file']['name']), 'r')) as f:
+            assert chan_info['data']['file']['format'] == 'raw', \
+                'Non-raw data format:\n%s' % pprint.pformat(chan_info)
+            dtype = self._zip_channel_dtype(chan_info)
+            data = numpy.frombuffer(
+                buffer(f.read()),
+                dtype=dtype,)
+        return data
+
+    def _zip_channel_dtype(self, chan_info):
+        type_ = chan_info['data']['type']
+        if type_ in ['float-data', 'float']:
+            dtype = numpy.dtype(numpy.float32)
+        elif type_ in ['integer-data', 'memory-integer-data']:
+            encoder = chan_info['data']['encoder']['type']
+            if encoder in ['signedinteger', 'signedinteger-limited']:
+                dtype = numpy.dtype(numpy.int32)
+            elif encoder in ['unsignedinteger', 'unsignedinteger-limited']:
+                dtype = numpy.dtype(numpy.uint32)
+            else:
+                raise ValueError('Unrecognized encoder type "%s" for "%s" data'
+                                 % (encoder, type_))
+        elif type_ in ['short-data', 'short', 'memory-short-data']:
+            encoder = chan_info['data']['encoder']['type']
+            if encoder in ['signedshort', 'signedshort-limited']:
+                dtype = numpy.dtype(numpy.int16)
+            elif encoder in ['unsignedshort', 'unsignedshort-limited']:
+                dtype = numpy.dtype(numpy.uint16)
+            else:
+                raise ValueError('Unrecognized encoder type "%s" for "%s" data'
+                                 % (encoder, type_))
+        else:
+            raise ValueError('Unrecognized data format "%s"' % type_)
+        byte_order = '>'
         # '>' (big endian) byte order.
         # From version 0.3 of JPKForceSpec.txt in the "Binary data" section:
         #    All forms of raw data are stored in chronological order
@@ -146,8 +172,7 @@ class JPKDriver (Driver):
         #    specified by the "channel.*.data.type" property, and is
         #    either short (2 bytes per value), integer (4 bytes), or
         #    float (4 bytes, IEEE format).
-        f.close()
-        return data
+        return dtype.newbyteorder(byte_order)
 
     def _zip_translate_params(self, params, chan_info):
         info = {