From 214ef941fcfe1e46e4d947145ca1e7649ee26bfa Mon Sep 17 00:00:00 2001
From: Jeff Reback <jeff@reback.net>
Date: Sun, 16 Oct 2016 11:44:32 -0400
Subject: [PATCH] ENH: add Pickle codec with support for object ndarrays

---
 .gitignore                       |  4 ++
 numcodecs/__init__.py            |  8 ++++
 numcodecs/msgpacks.py            | 52 ++++++++++++++++++++++++++
 numcodecs/pickles.py             | 64 ++++++++++++++++++++++++++++++++
 numcodecs/tests/common.py        | 33 +++++++++++++++-
 numcodecs/tests/test_msgpacks.py | 52 ++++++++++++++++++++++++++
 numcodecs/tests/test_pickle.py   | 48 ++++++++++++++++++++++++
 requirements.txt                 |  1 +
 8 files changed, 261 insertions(+), 1 deletion(-)
 create mode 100644 numcodecs/msgpacks.py
 create mode 100644 numcodecs/pickles.py
 create mode 100644 numcodecs/tests/test_msgpacks.py
 create mode 100644 numcodecs/tests/test_pickle.py

diff --git a/.gitignore b/.gitignore
index d2f47eb5..50aa2f21 100644
--- a/.gitignore
+++ b/.gitignore
@@ -6,6 +6,9 @@ __pycache__/
 # C extensions
 *.so
 
+# editor
+*~
+
 # Distribution / packaging
 .Python
 env/
@@ -44,6 +47,7 @@ nosetests.xml
 coverage.xml
 *,cover
 .hypothesis/
+cover/
 
 # Translations
 *.mo
diff --git a/numcodecs/__init__.py b/numcodecs/__init__.py
index 9a8c73e3..425d10df 100644
--- a/numcodecs/__init__.py
+++ b/numcodecs/__init__.py
@@ -61,6 +61,14 @@
 from numcodecs.categorize import Categorize
 register_codec(Categorize)
 
+from numcodecs.pickles import Pickle
+register_codec(Pickle)
+
+try:
+    from numcodecs.msgpacks import MsgPack
+    register_codec(MsgPack)
+except ImportError: # pragma: no cover
+    pass
 
 from numcodecs.checksum32 import CRC32, Adler32
 register_codec(CRC32)
diff --git a/numcodecs/msgpacks.py b/numcodecs/msgpacks.py
new file mode 100644
index 00000000..edddfd5e
--- /dev/null
+++ b/numcodecs/msgpacks.py
@@ -0,0 +1,52 @@
+# -*- coding: utf-8 -*-
+from __future__ import absolute_import, print_function, division
+
+
+import numpy as np
+
+
+from numcodecs.abc import Codec
+from numcodecs.compat import ndarray_from_buffer, buffer_copy
+import msgpack
+
+
+class MsgPack(Codec):
+    """Codec to encode data as msgpacked bytes. Useful for encoding python
+    strings
+
+    Raises
+    ------
+    encoding a non-object dtyped ndarray will raise ValueError
+
+    Examples
+    --------
+    >>> import numcodecs as codecs
+    >>> import numpy as np
+    >>> x = np.array(['foo', 'bar', 'baz'], dtype='object')
+    >>> f = codecs.MsgPack()
+    >>> f.decode(f.encode(x))
+    array(['foo', 'bar', 'baz'], dtype=object)
+
+    """  # flake8: noqa
+
+    codec_id = 'msgpack'
+
+    def encode(self, buf):
+        if hasattr(buf, 'dtype') and buf.dtype != 'object':
+            raise ValueError("cannot encode non-object ndarrays, %s "
+                             "dtype was passed" % buf.dtype)
+        return msgpack.packb(buf.tolist(), encoding='utf-8')
+
+    def decode(self, buf, out=None):
+        dec = np.array(msgpack.unpackb(buf, encoding='utf-8'), dtype='object')
+        if out is not None:
+            np.copyto(out, dec)
+            return out
+        else:
+            return dec
+
+    def get_config(self):
+        return dict(id=self.codec_id)
+
+    def __repr__(self):
+        return 'MsgPack()'
diff --git a/numcodecs/pickles.py b/numcodecs/pickles.py
new file mode 100644
index 00000000..d010c603
--- /dev/null
+++ b/numcodecs/pickles.py
@@ -0,0 +1,64 @@
+# -*- coding: utf-8 -*-
+from __future__ import absolute_import, print_function, division
+
+
+import numpy as np
+
+
+from numcodecs.abc import Codec
+from numcodecs.compat import ndarray_from_buffer, buffer_copy
+try:
+    import cPickle as pickle
+except ImportError:
+    import pickle
+
+
+class Pickle(Codec):
+    """Codec to encode data as as pickled bytes. Useful for encoding python
+    strings.
+
+    Parameters
+    ----------
+    protocol : int, defaults to pickle.HIGHEST_PROTOCOL
+        the protocol used to pickle data
+
+    Raises
+    ------
+    encoding a non-object dtyped ndarray will raise ValueError
+
+    Examples
+    --------
+    >>> import numcodecs as codecs
+    >>> import numpy as np
+    >>> x = np.array(['foo', 'bar', 'baz'], dtype='object')
+    >>> f = codecs.Pickle()
+    >>> f.decode(f.encode(x))
+    array(['foo', 'bar', 'baz'], dtype=object)
+
+    """  # flake8: noqa
+
+    codec_id = 'pickle'
+
+    def __init__(self, protocol=pickle.HIGHEST_PROTOCOL):
+        self.protocol = protocol
+
+    def encode(self, buf):
+        if hasattr(buf, 'dtype') and buf.dtype != 'object':
+            raise ValueError("cannot encode non-object ndarrays, %s "
+                             "dtype was passed" % buf.dtype)
+        return pickle.dumps(buf, protocol=self.protocol)
+
+    def decode(self, buf, out=None):
+        dec = pickle.loads(buf)
+        if out is not None:
+            np.copyto(out, dec)
+            return out
+        else:
+            return dec
+
+    def get_config(self):
+        return dict(id=self.codec_id,
+                    protocol=self.protocol)
+
+    def __repr__(self):
+        return 'Pickle(protocol=%s)' % self.protocol
diff --git a/numcodecs/tests/common.py b/numcodecs/tests/common.py
index 791de748..0089a781 100644
--- a/numcodecs/tests/common.py
+++ b/numcodecs/tests/common.py
@@ -5,7 +5,7 @@
 
 
 import numpy as np
-from nose.tools import eq_ as eq
+from nose.tools import eq_ as eq, assert_true
 from numpy.testing import assert_array_almost_equal
 
 
@@ -91,6 +91,37 @@ 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 compare(res, arr=arr):
+
+        assert_true(isinstance(res, np.ndarray))
+        assert_true(res.shape == arr.shape)
+        assert_true(res.dtype == 'object')
+
+        # numpy asserts don't compare object arrays
+        # properly; assert that we have the same nans
+        # and values
+        arr = arr.ravel().tolist()
+        res = res.ravel().tolist()
+        for a, r in zip(arr, res):
+            if a != a:
+                assert_true(r != r)
+            else:
+                assert_true(a == r)
+
+    enc = codec.encode(arr)
+    dec = codec.decode(enc)
+    compare(dec)
+
+    out = np.empty_like(arr)
+    codec.decode(enc, out=out)
+    compare(out)
+
+
 def check_config(codec):
     config = codec.get_config()
     # round-trip through JSON to check serialization
diff --git a/numcodecs/tests/test_msgpacks.py b/numcodecs/tests/test_msgpacks.py
new file mode 100644
index 00000000..1c5672d2
--- /dev/null
+++ b/numcodecs/tests/test_msgpacks.py
@@ -0,0 +1,52 @@
+# -*- coding: utf-8 -*-
+from __future__ import absolute_import, print_function, division
+
+
+import nose
+import numpy as np
+from numpy.testing import assert_raises
+
+try:
+    from numcodecs.msgpacks import MsgPack
+except ImportError:
+    raise nose.SkipTest("no msgpack installed")
+
+from numcodecs.tests.common import (check_config, check_repr,
+                                    check_encode_decode_objects)
+
+
+# object array with strings
+# object array with mix strings / nans
+# object array with mix of string, int, float
+arrays = [
+    np.array(['foo', 'bar', 'baz'] * 300, dtype=object),
+    np.array([['foo', 'bar', np.nan]] * 300, dtype=object),
+    np.array(['foo', 1.0, 2] * 300, dtype=object),
+]
+
+# non-object ndarrays
+arrays_incompat = [
+    np.arange(1000, dtype='i4'),
+    np.array(['foo', 'bar', 'baz'] * 300),
+]
+
+
+def test_encode_errors():
+    for arr in arrays_incompat:
+        codec = MsgPack()
+        assert_raises(ValueError, codec.encode, arr)
+
+
+def test_encode_decode():
+    for arr in arrays:
+        codec = MsgPack()
+        check_encode_decode_objects(arr, codec)
+
+
+def test_config():
+    codec = MsgPack()
+    check_config(codec)
+
+
+def test_repr():
+    check_repr("MsgPack()")
diff --git a/numcodecs/tests/test_pickle.py b/numcodecs/tests/test_pickle.py
new file mode 100644
index 00000000..dcad15b3
--- /dev/null
+++ b/numcodecs/tests/test_pickle.py
@@ -0,0 +1,48 @@
+# -*- coding: utf-8 -*-
+from __future__ import absolute_import, print_function, division
+
+
+import numpy as np
+from numpy.testing import assert_raises
+
+
+from numcodecs.pickles import Pickle
+from numcodecs.tests.common import (check_config, check_repr,
+                                    check_encode_decode_objects)
+
+
+# object array with strings
+# object array with mix strings / nans
+# object array with mix of string, int, float
+arrays = [
+    np.array(['foo', 'bar', 'baz'] * 300, dtype=object),
+    np.array([['foo', 'bar', np.nan]] * 300, dtype=object),
+    np.array(['foo', 1.0, 2] * 300, dtype=object),
+]
+
+# non-object ndarrays
+arrays_incompat = [
+    np.arange(1000, dtype='i4'),
+    np.array(['foo', 'bar', 'baz'] * 300),
+]
+
+
+def test_encode_errors():
+    for arr in arrays_incompat:
+        codec = Pickle()
+        assert_raises(ValueError, codec.encode, arr)
+
+
+def test_encode_decode():
+    for arr in arrays:
+        codec = Pickle()
+        check_encode_decode_objects(arr, codec)
+
+
+def test_config():
+    codec = Pickle(protocol=-1)
+    check_config(codec)
+
+
+def test_repr():
+    check_repr("Pickle(protocol=-1)")
diff --git a/requirements.txt b/requirements.txt
index 24ce15ab..82a67644 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -1 +1,2 @@
 numpy
+msgpack-python