Merge pull request #5319 from larsmans/c-fixes

MAINT: refactor packbits/unpackbits

1 files changed, 26 insertions, 0 deletions

diff --git a/numpy/lib/tests/test_packbits.py b/numpy/lib/tests/test_packbits.py
new file mode 100644
index 000000000..186e8960d
--- /dev/null
+++ b/numpy/lib/tests/test_packbits.py
@@ -0,0 +1,26 @@
+import numpy as np
+
+from numpy.testing import assert_array_equal, assert_equal, assert_raises
+
+
+def test_packbits():
+    # Copied from the docstring.
+    a = [[[1, 0, 1], [0, 1, 0]],
+         [[1, 1, 0], [0, 0, 1]]]
+    for dtype in [np.bool, np.uint8, np.int]:
+        arr = np.array(a, dtype=dtype)
+        b = np.packbits(arr, axis=-1)
+        assert_equal(b.dtype, np.uint8)
+        assert_array_equal(b, np.array([[[160], [64]], [[192], [32]]]))
+
+    assert_raises(TypeError, np.packbits, np.array(a, dtype=float))
+
+
+def test_unpackbits():
+    # Copied from the docstring.
+    a = np.array([[2], [7], [23]], dtype=np.uint8)
+    b = np.unpackbits(a, axis=1)
+    assert_equal(b.dtype, np.uint8)
+    assert_array_equal(b, np.array([[0, 0, 0, 0, 0, 0, 1, 0],
+                                    [0, 0, 0, 0, 0, 1, 1, 1],
+                                    [0, 0, 0, 1, 0, 1, 1, 1]]))