Additional dtypes for labels. (#197)

numbagg/decorators.py

@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 import abc
+import itertools
 from collections.abc import Iterable
 from functools import cache, cached_property
 from typing import Any, Callable, TypeVar
@@ -378,21 +379,17 @@ def __init__(
         self.func = func
 
         if signature is None:
-            signature = [
-                (numba.float32, numba.int32, numba.float32),
-                (numba.float32, numba.int64, numba.float32),
-                (numba.float64, numba.int32, numba.float64),
-                (numba.float64, numba.int64, numba.float64),
-            ]
-
+            values_dtypes: tuple[numba.dtype, ...] = (numba.float32, numba.float64)
+            labels_dtypes = (numba.int8, numba.int16, numba.int32, numba.int64)
             if supports_ints:
-                signature += [
-                    (numba.int32, numba.int32, numba.int32),
-                    (numba.int32, numba.int64, numba.int32),
-                    (numba.int64, numba.int32, numba.int64),
-                    (numba.int64, numba.int64, numba.int64),
-                ]
+                values_dtypes += (numba.int32, numba.int64)
 
+            signature = [
+                (value_type, label_type, value_type)
+                for value_type, label_type in itertools.product(
+                    values_dtypes, labels_dtypes
+                )
+            ]
         for sig in signature:
             if not isinstance(sig, tuple):
                 raise TypeError(f"signatures for ndmoving must be tuples: {signature}")

numbagg/test/test_grouped.py

@@ -222,13 +222,14 @@ def test_additional_dim_equivalence(func, values, labels, dtype):
     assert_almost_equal(result, expected)
 
 
+@pytest.mark.parametrize("labels_type", [np.int8, np.int16, np.int32, np.int64])
 @pytest.mark.parametrize("func, _, npfunc", [f for f in FUNCTIONS_CONSTANT])
-def test_group_func_axis_1d_labels(func, _, npfunc):
+def test_group_func_axis_1d_labels(func, _, npfunc, labels_type):
     if npfunc is None:
         pytest.skip("No numpy equivalent")
 
     values = np.arange(5.0)
-    labels = np.arange(5)
+    labels = np.arange(5, dtype=labels_type)
     result = func(values, labels)
     assert_almost_equal(result, values)