From cc7e4c1a703fded4964f7e6a0f1dfec98e3d1cca Mon Sep 17 00:00:00 2001 From: Dmitry Kalinkin Date: Sat, 30 Mar 2024 02:13:05 -0400 Subject: [PATCH] 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 --------- Signed-off-by: Henry Schreiner Co-authored-by: Henry Schreiner --- include/bh_python/fill.hpp | 4 +++- tests/test_histogram.py | 11 +++++++++++ 2 files changed, 14 insertions(+), 1 deletion(-) diff --git a/include/bh_python/fill.hpp b/include/bh_python/fill.hpp index 9e630db5..fbd36b43 100644 --- a/include/bh_python/fill.hpp +++ b/include/bh_python/fill.hpp @@ -102,7 +102,9 @@ inline decltype(auto) special_cast>(py::handle x) { auto dtype = py::cast(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>(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 index 9c039d4e..da3a95e1 100644 --- 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)