fix: allow filling Int axis with unsigned integers (#917)

* fix: allow filling Int axis with unsigned integers * tests: add test for uint filling Signed-off-by: Henry Schreiner <[email protected]> --------- Signed-off-by: Henry Schreiner <[email protected]> Co-authored-by: Henry Schreiner <[email protected]>
scikit-hep · Mar 30, 2024 · cc7e4c1 · cc7e4c1
1 parent e8e2487
commit cc7e4c1
Show file tree

Hide file tree

Showing 2 changed files with 14 additions and 1 deletion.
diff --git a/include/bh_python/fill.hpp b/include/bh_python/fill.hpp
@@ -102,7 +102,9 @@ inline decltype(auto) special_cast<c_array_t<int>>(py::handle x) {
     auto dtype = py::cast<py::array>(x).dtype();
     if(dtype.equal(np.attr("bool_")) || dtype.equal(np.attr("int8"))
        || dtype.equal(np.attr("int16")) || dtype.equal(np.attr("int32"))
-       || dtype.equal(np.attr("int64")))
+       || dtype.equal(np.attr("int64")) || dtype.equal(np.attr("uint8"))
+       || dtype.equal(np.attr("uint16")) || dtype.equal(np.attr("uint32"))
+       || dtype.equal(np.attr("uint64")))
         return py::cast<c_array_t<int>>(x);
     throw py::type_error("Only integer arrays supported when targeting integer axes");
 }

diff --git a/tests/test_histogram.py b/tests/test_histogram.py
@@ -939,6 +939,17 @@ def test_numpy_conversion_5():
     assert a1[2, 1] == 5
 
 
+@pytest.mark.parametrize(
+    "dtype",
+    [np.int8, np.int16, np.int32, np.int64, np.uint8, np.uint16, np.uint32, np.uint64],
+)
+def test_fill_dtypes(dtype):
+    a = bh.Histogram(bh.axis.Integer(0, 2), storage=bh.storage.Int64())
+    a.fill(np.array([0, 0, 0, 1, 1, 2], dtype=dtype))
+    a.fill(dtype(0))
+    assert list(a.values()) == [4, 2]
+
+
 def test_fill_with_sequence_0():
     def fa(*args):
         return np.array(args, dtype=float)