generalise pickle and msgpack to handle arrays with any dtype

Author: alimanfoo

numcodecs/msgpacks.py

@@ -10,7 +10,8 @@
 
 
 class MsgPack(Codec):
-    """Codec to encode data as msgpacked bytes. Useful for encoding an array of Python strings.
+    """Codec to encode data as msgpacked bytes. Useful for encoding an array of Python string
+    objects.
 
     Examples
     --------
@@ -33,20 +34,27 @@ class MsgPack(Codec):
 
     codec_id = 'msgpack'
 
+    def __init__(self, encoding='utf-8'):
+        self.encoding = encoding
+
     def encode(self, buf):
-        buf = np.asarray(buf, dtype='object')
-        return msgpack.packb(buf.tolist(), encoding='utf-8')
+        buf = np.asarray(buf)
+        l = buf.tolist()
+        l.append(buf.dtype.str)
+        return msgpack.packb(l, encoding=self.encoding)
 
     def decode(self, buf, out=None):
-        dec = np.array(msgpack.unpackb(buf, encoding='utf-8'), dtype='object')
+        l = msgpack.unpackb(buf, encoding=self.encoding)
+        dec = np.array(l[:-1], dtype=l[-1])
         if out is not None:
             np.copyto(out, dec)
             return out
         else:
             return dec
 
     def get_config(self):
-        return dict(id=self.codec_id)
+        return dict(id=self.codec_id,
+                    encoding=self.encoding)
 
     def __repr__(self):
-        return 'MsgPack()'
+        return 'MsgPack(encoding=%r)' % self.encoding

numcodecs/pickles.py

@@ -14,7 +14,8 @@
 
 
 class Pickle(Codec):
-    """Codec to encode data as as pickled bytes. Useful for encoding an array of Python strings.
+    """Codec to encode data as as pickled bytes. Useful for encoding an array of Python string
+    objects.
 
     Parameters
     ----------
@@ -42,7 +43,6 @@ def __init__(self, protocol=pickle.HIGHEST_PROTOCOL):
         self.protocol = protocol
 
     def encode(self, buf):
-        buf = np.asarray(buf, dtype='object')
         return pickle.dumps(buf, protocol=self.protocol)
 
     def decode(self, buf, out=None):

numcodecs/tests/common.py

@@ -91,16 +91,13 @@ def compare(res):
     compare(out)
 
 
-def check_encode_decode_objects(arr, codec):
-
-    # this is a more specific test that check_encode_decode
-    # as these require actual objects (and not bytes only)
+def check_encode_decode_array(arr, codec):
 
     def compare(res, arr=arr):
 
         assert_true(isinstance(res, np.ndarray))
         assert_true(res.shape == arr.shape)
-        assert_true(res.dtype == 'object')
+        assert_true(res.dtype == arr.dtype)
 
         # numpy asserts don't compare object arrays
         # properly; assert that we have the same nans
@@ -117,7 +114,7 @@ def compare(res, arr=arr):
     dec = codec.decode(enc)
     compare(dec)
 
-    out = np.empty_like(arr, dtype='object')
+    out = np.empty_like(arr)
     codec.decode(enc, out=out)
     compare(out)
 

numcodecs/tests/test_msgpacks.py

@@ -10,7 +10,7 @@
 except ImportError:  # pragma: no cover
     raise nose.SkipTest("msgpack-python not available")
 
-from numcodecs.tests.common import check_config, check_repr, check_encode_decode_objects
+from numcodecs.tests.common import check_config, check_repr, check_encode_decode_array
 
 
 # object array with strings
@@ -28,7 +28,7 @@
 def test_encode_decode():
     for arr in arrays:
         codec = MsgPack()
-        check_encode_decode_objects(arr, codec)
+        check_encode_decode_array(arr, codec)
 
 
 def test_config():
@@ -37,4 +37,5 @@ def test_config():
 
 
 def test_repr():
-    check_repr("MsgPack()")
+    check_repr("MsgPack(encoding='utf-8')")
+    check_repr("MsgPack(encoding='ascii')")

numcodecs/tests/test_pickles.py

@@ -6,7 +6,7 @@
 
 
 from numcodecs.pickles import Pickle
-from numcodecs.tests.common import check_config, check_repr, check_encode_decode_objects
+from numcodecs.tests.common import check_config, check_repr, check_encode_decode_array
 
 
 # object array with strings
@@ -24,7 +24,7 @@
 def test_encode_decode():
     codec = Pickle()
     for arr in arrays:
-        check_encode_decode_objects(arr, codec)
+        check_encode_decode_array(arr, codec)
 
 
 def test_config():