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..dcad5c6be 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 - +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,68 @@ 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, - numpy.int_, - numpy.float_, - numpy.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. - - :param numpy.ndarray array: Array of string identifiers, or of enum - items, to encode. + Encode an array of enum items or string identifiers into an EnumArray. - :returns: An :any:`EnumArray` encoding the input array values. - :rtype: :any:`EnumArray` + Args: + array: The input array to encode. Can be an EnumArray, a NumPy array + of enum items, or a NumPy array of string identifiers. - For instance: + Returns: + An EnumArray containing the encoded values. - >>> string_identifier_array = asarray(['free_lodger', 'owner']) - >>> encoded_array = HousingOccupancyStatus.encode(string_identifier_array) - >>> encoded_array[0] - 2 # Encoded value + 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) - >>> 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, numpy.ndarray) and array.dtype.kind in {"U", "S"}: - array = numpy.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, numpy.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 = numpy.select( - [array == item for item in cls], - [item.index for item in cls], - ).astype(ENUM_ARRAY_DTYPE) - else: - array = numpy.select( - [array.astype(str) == item.name for item in cls], - [item.index for item in cls], - ).astype(ENUM_ARRAY_DTYPE) - - return EnumArray(array, cls) + ) + 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], + ) + elif array.dtype.kind in {"i", "u"}: + # Integer array + indices = array + else: + raise ValueError(f"Unsupported array dtype: {array.dtype}") + + return EnumArray(indices.astype(ENUM_ARRAY_DTYPE), cls)