From 4fec397e0a4cbff37ba9cf553fe92d0c8662891a Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Tue, 23 Apr 2024 01:45:34 -0700 Subject: [PATCH 1/3] FutureWarning regarding bool checking in enum.py Fixes #100 --- changelog_entry.yaml | 4 ++++ policyengine_core/enums/enum.py | 18 +++++++++--------- 2 files changed, 13 insertions(+), 9 deletions(-) diff --git a/changelog_entry.yaml b/changelog_entry.yaml index e69de29bb..89ea53751 100644 --- a/changelog_entry.yaml +++ b/changelog_entry.yaml @@ -0,0 +1,4 @@ +- bump: patch + changes: + fixed: + - FutureWarning issue with bools and enums. diff --git a/policyengine_core/enums/enum.py b/policyengine_core/enums/enum.py index a6814a379..9e311c824 100644 --- a/policyengine_core/enums/enum.py +++ b/policyengine_core/enums/enum.py @@ -3,7 +3,7 @@ import enum from typing import Union -import numpy +import numpy as np from .config import ENUM_ARRAY_DTYPE from .enum_array import EnumArray @@ -38,9 +38,9 @@ def encode( cls, array: Union[ EnumArray, - numpy.int_, - numpy.float_, - numpy.object_, + np.int_, + np.float_, + np.object_, ], ) -> EnumArray: """ @@ -75,14 +75,14 @@ def encode( array = array.astype(str) # String array - if isinstance(array, numpy.ndarray) and array.dtype.kind in {"U", "S"}: - array = numpy.select( + if isinstance(array, np.ndarray) and array.dtype.kind in {"U", "S"}: + array = np.select( [array == item.name for item in cls], [item.index for item in cls], ).astype(ENUM_ARRAY_DTYPE) # Enum items arrays - elif isinstance(array, numpy.ndarray) and array.dtype.kind == "O": + elif isinstance(array, np.ndarray) and array.dtype.kind == "O": # Ensure we are comparing the comparable. The problem this fixes: # On entering this method "cls" will generally come from # variable.possible_values, while the array values may come from @@ -97,12 +97,12 @@ def encode( if len(array) > 0 and cls.__name__ is array[0].__class__.__name__: cls = array[0].__class__ if array[0].__class__.__name__ != "bytes": - array = numpy.select( + array = np.select( [array == item for item in cls], [item.index for item in cls], ).astype(ENUM_ARRAY_DTYPE) else: - array = numpy.select( + array = np.select( [array.astype(str) == item.name for item in cls], [item.index for item in cls], ).astype(ENUM_ARRAY_DTYPE) From 6e9f96233118fdef0d86df12886781bea97d0540 Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Tue, 23 Apr 2024 01:53:15 -0700 Subject: [PATCH 2/3] refactor --- policyengine_core/enums/enum.py | 116 +++++++++++--------------------- 1 file changed, 41 insertions(+), 75 deletions(-) diff --git a/policyengine_core/enums/enum.py b/policyengine_core/enums/enum.py index 9e311c824..2141c94c8 100644 --- a/policyengine_core/enums/enum.py +++ b/policyengine_core/enums/enum.py @@ -1,17 +1,10 @@ from __future__ import annotations - import enum from typing import Union - import numpy as np - from .config import ENUM_ARRAY_DTYPE from .enum_array import EnumArray -import warnings - -warnings.simplefilter("ignore", category=FutureWarning) - class Enum(enum.Enum): """ @@ -19,92 +12,65 @@ class Enum(enum.Enum): have an index. """ - # Tweak enums to add an index attribute to each enum item def __init__(self, name: str) -> None: - # When the enum item is initialized, self._member_names_ contains the - # names of the previously initialized items, so its length is the index - # of this item. + """ + Initialize an Enum item with a name and an index. + + The index is automatically assigned based on the order of the Enum items. + """ self.index = len(self._member_names_) - # Bypass the slow Enum.__eq__ __eq__ = object.__eq__ - - # In Python 3, __hash__ must be defined if __eq__ is defined to stay - # hashable. __hash__ = object.__hash__ @classmethod - def encode( - cls, - array: Union[ - EnumArray, - np.int_, - np.float_, - np.object_, - ], - ) -> EnumArray: + def encode(cls, array: Union[EnumArray, np.ndarray]) -> EnumArray: """ - Encode a string numpy array, an enum item numpy array, or an int numpy - array into an :any:`EnumArray`. See :any:`EnumArray.decode` for - decoding. + Encode an array of enum items or string identifiers into an EnumArray. - :param numpy.ndarray array: Array of string identifiers, or of enum - items, to encode. + Args: + array: The input array to encode. Can be an EnumArray, a NumPy array + of enum items, or a NumPy array of string identifiers. - :returns: An :any:`EnumArray` encoding the input array values. - :rtype: :any:`EnumArray` + Returns: + An EnumArray containing the encoded values. - For instance: + Examples: + >>> string_array = np.array(["ITEM_1", "ITEM_2", "ITEM_3"]) + >>> encoded_array = MyEnum.encode(string_array) + >>> encoded_array + EnumArray([1, 2, 3], dtype=int8) - >>> string_identifier_array = asarray(['free_lodger', 'owner']) - >>> encoded_array = HousingOccupancyStatus.encode(string_identifier_array) - >>> encoded_array[0] - 2 # Encoded value - - >>> free_lodger = HousingOccupancyStatus.free_lodger - >>> owner = HousingOccupancyStatus.owner - >>> enum_item_array = asarray([free_lodger, owner]) - >>> encoded_array = HousingOccupancyStatus.encode(enum_item_array) - >>> encoded_array[0] - 2 # Encoded value + >>> item_array = np.array([MyEnum.ITEM_1, MyEnum.ITEM_2, MyEnum.ITEM_3]) + >>> encoded_array = MyEnum.encode(item_array) + >>> encoded_array + EnumArray([1, 2, 3], dtype=int8) """ if isinstance(array, EnumArray): return array - if isinstance(array == 0, bool): + if array.dtype.kind == "b": + # Convert boolean array to string array array = array.astype(str) - # String array - if isinstance(array, np.ndarray) and array.dtype.kind in {"U", "S"}: - array = np.select( + if array.dtype.kind in {"U", "S"}: + # String array + indices = np.select( [array == item.name for item in cls], [item.index for item in cls], - ).astype(ENUM_ARRAY_DTYPE) - - # Enum items arrays - elif isinstance(array, np.ndarray) and array.dtype.kind == "O": - # Ensure we are comparing the comparable. The problem this fixes: - # On entering this method "cls" will generally come from - # variable.possible_values, while the array values may come from - # directly importing a module containing an Enum class. However, - # variables (and hence their possible_values) are loaded by a call - # to load_module, which gives them a different identity from the - # ones imported in the usual way. - # - # So, instead of relying on the "cls" passed in, we use only its - # name to check that the values in the array, if non-empty, are of - # the right type. - if len(array) > 0 and cls.__name__ is array[0].__class__.__name__: - cls = array[0].__class__ - if array[0].__class__.__name__ != "bytes": - array = np.select( - [array == item for item in cls], - [item.index for item in cls], - ).astype(ENUM_ARRAY_DTYPE) - else: - array = np.select( - [array.astype(str) == item.name for item in cls], - [item.index for item in cls], - ).astype(ENUM_ARRAY_DTYPE) + ) + elif array.dtype.kind == "O": + # Enum items array + if len(array) > 0: + first_item = array[0] + if cls.__name__ == type(first_item).__name__: + # Use the same Enum class as the array items + cls = type(first_item) + indices = np.select( + [array == item for item in cls], + [item.index for item in cls], + ) + else: + raise ValueError(f"Unsupported array dtype: {array.dtype}") - return EnumArray(array, cls) + return EnumArray(indices.astype(ENUM_ARRAY_DTYPE), cls) From 6803f6eb48f3fdf5ec4cef55c30f1d7e7d52938d Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Fri, 26 Apr 2024 23:45:12 -0700 Subject: [PATCH 3/3] 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 2141c94c8..dcad5c6be 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}")