ENH ensure np.packbits works on np.bool dtype

numpy · Nov 30, 2014 · 24effb6 · 24effb6
1 parent 6ce9883
commit 24effb6
Show file tree

Hide file tree

Showing 2 changed files with 12 additions and 8 deletions.
diff --git a/numpy/lib/src/_compiled_base.c b/numpy/lib/src/_compiled_base.c
@@ -1376,9 +1376,9 @@ pack_bits(PyObject *input, int axis)
     if (inp == NULL) {
         return NULL;
     }
-    if (!PyArray_ISINTEGER(inp)) {
+    if (!PyArray_ISBOOL(inp) && !PyArray_ISINTEGER(inp)) {
         PyErr_SetString(PyExc_TypeError,
-                "Expected an input array of integer data type");
+                "Expected an input array of integer or boolean data type");
         goto fail;
     }
 

diff --git a/numpy/lib/tests/test_packbits.py b/numpy/lib/tests/test_packbits.py
@@ -1,15 +1,19 @@
 import numpy as np
 
-from numpy.testing import assert_array_equal, assert_equal
+from numpy.testing import assert_array_equal, assert_equal, assert_raises
 
 
 def test_packbits():
     # Copied from the docstring.
-    a = np.array([[[1, 0, 1], [0, 1, 0]],
-                  [[1, 1, 0], [0, 0, 1]]])
-    b = np.packbits(a, axis=-1)
-    assert_equal(b.dtype, np.uint8)
-    assert_array_equal(b, np.array([[[160], [64]], [[192], [32]]]))
+    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():