From 6803f6eb48f3fdf5ec4cef55c30f1d7e7d52938d Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Fri, 26 Apr 2024 23:45:12 -0700 Subject: [PATCH] Added a new elif condition to handle integer arrays. If the array's data type kind is "i" (signed integer) or "u" (unsigned integer), the indices variable is directly assigned the array value. Updated the else condition to raise a ValueError only for unsupported array data types that are not handled by the previous conditions. --- policyengine_core/enums/enum.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/policyengine_core/enums/enum.py b/policyengine_core/enums/enum.py index 2141c94c..dcad5c6b 100644 --- a/policyengine_core/enums/enum.py +++ b/policyengine_core/enums/enum.py @@ -70,6 +70,9 @@ def encode(cls, array: Union[EnumArray, np.ndarray]) -> EnumArray: [array == item for item in cls], [item.index for item in cls], ) + elif array.dtype.kind in {"i", "u"}: + # Integer array + indices = array else: raise ValueError(f"Unsupported array dtype: {array.dtype}")