fix: Replace numpy-Python comparison with dtype

PolicyEngine · May 8, 2024 · bf0a54d · bf0a54d
1 parent ac7317c
commit bf0a54d
Show file tree

Hide file tree

Showing 2 changed files with 9 additions and 2 deletions.
diff --git a/changelog_entry.yaml b/changelog_entry.yaml
@@ -0,0 +1,4 @@
+- bump: patch
+  changes:
+    changed:
+    - Replaced unsafe numpy-Python comparison with use of numpy dtype to convert byte-string arrays to Unicode ones within enums
diff --git a/policyengine_core/enums/enum.py b/policyengine_core/enums/enum.py
@@ -49,8 +49,11 @@ def encode(cls, array: Union[EnumArray, np.ndarray]) -> EnumArray:
         if isinstance(array, EnumArray):
             return array
 
-        # if array.dtype.kind == "b":
-        if isinstance(array == 0, bool):
+        # First, convert byte-string arrays to Unicode-string arrays
+        # Confusingly, Numpy uses "S" to refer to byte-string arrays
+        # and "U" to refer to Unicode-string arrays, which are also
+        # referred to as the "str" type
+        if array.dtype.kind == "S":
             # Convert boolean array to string array
             array = array.astype(str)