fix(enums): do actual indexing (#1267)

Author: bonjourmauko

openfisca_core/indexed_enums/__init__.py

@@ -1,6 +1,7 @@
 """Enumerations for variables with a limited set of possible values."""
 
 from . import types
+from ._enum_type import EnumType
 from .config import ENUM_ARRAY_DTYPE
 from .enum import Enum
 from .enum_array import EnumArray
@@ -9,5 +10,6 @@
     "ENUM_ARRAY_DTYPE",
     "Enum",
     "EnumArray",
+    "EnumType",
     "types",
 ]

openfisca_core/indexed_enums/_enum_type.py

@@ -0,0 +1,113 @@
+from __future__ import annotations
+
+from typing import final
+
+import numpy
+
+from . import types as t
+
+
+def _item_list(enum_class: type[t.Enum]) -> t.ItemList:
+    """Return the non-vectorised list of enum items."""
+    return [
+        (index, name, value)
+        for index, (name, value) in enumerate(enum_class.__members__.items())
+    ]
+
+
+def _item_dtype(enum_class: type[t.Enum]) -> t.RecDType:
+    """Return the dtype of the indexed enum's items."""
+    size = max(map(len, enum_class.__members__.keys()))
+    return numpy.dtype(
+        (
+            numpy.generic,
+            {
+                "index": (t.EnumDType, 0),
+                "name": (f"U{size}", 2),
+                "enum": (enum_class, 2 + size * 4),
+            },
+        )
+    )
+
+
+def _item_array(enum_class: type[t.Enum]) -> t.RecArray:
+    """Return the indexed enum's items."""
+    items = _item_list(enum_class)
+    dtype = _item_dtype(enum_class)
+    array = numpy.array(items, dtype=dtype)
+    return array.view(numpy.recarray)
+
+
+@final
+class EnumType(t.EnumType):
+    """Meta class for creating an indexed :class:`.Enum`.
+
+    Examples:
+        >>> from openfisca_core import indexed_enums as enum
+
+        >>> class Enum(enum.Enum, metaclass=enum.EnumType):
+        ...     pass
+
+        >>> Enum.items
+        Traceback (most recent call last):
+        AttributeError: type object 'Enum' has no attribute 'items'
+
+        >>> class Housing(Enum):
+        ...     OWNER = "Owner"
+        ...     TENANT = "Tenant"
+
+        >>> Housing.items
+        rec.array([(0, 'OWNER', <Housing.OWNER: 'Owner'>), ...])
+
+        >>> Housing.indices
+        array([0, 1], dtype=int16)
+
+        >>> Housing.names
+        array(['OWNER', 'TENANT'], dtype='<U6')
+
+        >>> Housing.enums
+        array([<Housing.OWNER: 'Owner'>, <Housing.TENANT: 'Tenant'>], dtype...)
+
+    """
+
+    #: The items of the indexed enum class.
+    items: t.RecArray
+
+    @property
+    def indices(cls) -> t.IndexArray:
+        """Return the indices of the indexed enum class."""
+        return cls.items.index
+
+    @property
+    def names(cls) -> t.StrArray:
+        """Return the names of the indexed enum class."""
+        return cls.items.name
+
+    @property
+    def enums(cls) -> t.ObjArray:
+        """Return the members of the indexed enum class."""
+        return cls.items.enum
+
+    def __new__(
+        metacls,
+        cls: str,
+        bases: tuple[type, ...],
+        classdict: t.EnumDict,
+        **kwds: object,
+    ) -> t.EnumType:
+        """Create a new indexed enum class."""
+        # Create the enum class.
+        enum_class = super().__new__(metacls, cls, bases, classdict, **kwds)
+
+        # If the enum class has no members, return it as is.
+        if not enum_class.__members__:
+            return enum_class
+
+        # Add the items attribute to the enum class.
+        enum_class.items = _item_array(enum_class)
+
+        # Return the modified enum class.
+        return enum_class
+
+    def __dir__(cls) -> list[str]:
+        return sorted({"items", "indices", "names", "enums", *super().__dir__()})

openfisca_core/indexed_enums/_type_guards.py

@@ -0,0 +1,68 @@
+from __future__ import annotations
+
+from typing_extensions import TypeIs
+
+import numpy
+
+from . import types as t
+
+
+def _is_int_array(array: t.AnyArray) -> TypeIs[t.IndexArray | t.IntArray]:
+    """Narrow the type of a given array to an array of :obj:`numpy.integer`.
+
+    Args:
+        array: Array to check.
+
+    Returns:
+        bool: True if ``array`` is an array of :obj:`numpy.integer`, False otherwise.
+
+    Examples:
+        >>> import numpy
+
+        >>> array = numpy.array([1], dtype=numpy.int16)
+        >>> _is_int_array(array)
+        True
+
+        >>> array = numpy.array([1], dtype=numpy.int32)
+        >>> _is_int_array(array)
+        True
+
+        >>> array = numpy.array([1.0])
+        >>> _is_int_array(array)
+        False
+
+    """
+    return numpy.issubdtype(array.dtype, numpy.integer)
+
+
+def _is_str_array(array: t.AnyArray) -> TypeIs[t.StrArray]:
+    """Narrow the type of a given array to an array of :obj:`numpy.str_`.
+
+    Args:
+        array: Array to check.
+
+    Returns:
+        bool: True if ``array`` is an array of :obj:`numpy.str_`, False otherwise.
+
+    Examples:
+        >>> import numpy
+
+        >>> from openfisca_core import indexed_enums as enum
+
+        >>> class Housing(enum.Enum):
+        ...     OWNER = "owner"
+        ...     TENANT = "tenant"
+
+        >>> array = numpy.array([Housing.OWNER])
+        >>> _is_str_array(array)
+        False
+
+        >>> array = numpy.array(["owner"])
+        >>> _is_str_array(array)
+        True
+
+    """
+    return numpy.issubdtype(array.dtype, str)
+
+
+__all__ = ["_is_int_array", "_is_str_array"]

openfisca_core/indexed_enums/enum.py

@@ -3,11 +3,12 @@
 import numpy
 
 from . import types as t
-from .config import ENUM_ARRAY_DTYPE
+from ._enum_type import EnumType
+from ._type_guards import _is_int_array, _is_str_array
 from .enum_array import EnumArray
 
 
-class Enum(t.Enum):
+class Enum(t.Enum, metaclass=EnumType):
     """Enum based on `enum34 <https://pypi.python.org/pypi/enum34/>`_.
 
     Its items have an :class:`int` index, useful and performant when running
@@ -115,20 +116,19 @@ def __ne__(self, other: object) -> bool:
             return NotImplemented
         return self.index != other.index
 
-    #: :meth:`.__hash__` must also be defined so as to stay hashable.
-    __hash__ = object.__hash__
+    def __hash__(self) -> int:
+        return hash(self.index)
 
     @classmethod
     def encode(
         cls,
         array: (
             EnumArray
-            | t.Array[t.DTypeStr]
-            | t.Array[t.DTypeInt]
-            | t.Array[t.DTypeEnum]
-            | t.Array[t.DTypeObject]
-            | t.ArrayLike[str]
+            | t.IntArray
+            | t.StrArray
+            | t.ObjArray
             | t.ArrayLike[int]
+            | t.ArrayLike[str]
             | t.ArrayLike[t.Enum]
         ),
     ) -> EnumArray:
@@ -143,7 +143,6 @@ def encode(
         Raises:
             TypeError: If ``array`` is a scalar :class:`~numpy.ndarray`.
             TypeError: If ``array`` is of a diffent :class:`.Enum` type.
-            NotImplementedError: If ``array`` is of an unsupported type.
 
         Examples:
             >>> import numpy
@@ -187,7 +186,7 @@ def encode(
             >>> array = numpy.array([b"TENANT"])
             >>> enum_array = Housing.encode(array)
             Traceback (most recent call last):
-            NotImplementedError: Unsupported encoding: bytes48.
+            TypeError: Failed to encode "[b'TENANT']" of type 'bytes_', as i...
 
         .. seealso::
             :meth:`.EnumArray.decode` for decoding.
@@ -200,7 +199,7 @@ def encode(
             return cls.encode(numpy.array(array))
 
         if array.size == 0:
-            return EnumArray(array, cls)
+            return EnumArray(numpy.array([]), cls)
 
         if array.ndim == 0:
             msg = (
@@ -209,49 +208,37 @@ def encode(
             )
             raise TypeError(msg)
 
-        # Enum data type array
-        if numpy.issubdtype(array.dtype, t.DTypeEnum):
-            indexes = numpy.array([item.index for item in cls], t.DTypeEnum)
-            return EnumArray(indexes[array[array < indexes.size]], cls)
-
         # Integer array
-        if numpy.issubdtype(array.dtype, int):
-            array = numpy.array(array, dtype=t.DTypeEnum)
-            return cls.encode(array)
+        if _is_int_array(array):
+            indices = numpy.array(array[array < len(cls.items)], dtype=t.EnumDType)
+            return EnumArray(indices, cls)
 
         # String array
-        if numpy.issubdtype(array.dtype, t.DTypeStr):
-            enums = [cls.__members__[key] for key in array if key in cls.__members__]
-            return cls.encode(enums)
-
-        # Enum items arrays
-        if numpy.issubdtype(array.dtype, t.DTypeObject):
-            # 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 cls.__name__ is array[0].__class__.__name__:
-                array = numpy.select(
-                    [array == item for item in array[0].__class__],
-                    [item.index for item in array[0].__class__],
-                ).astype(ENUM_ARRAY_DTYPE)
-                return EnumArray(array, cls)
-
-            msg = (
-                f"Diverging enum types are not supported: expected {cls.__name__}, "
-                f"but got {array[0].__class__.__name__} instead."
-            )
-            raise TypeError(msg)
-
-        msg = f"Unsupported encoding: {array.dtype.name}."
-        raise NotImplementedError(msg)
+        if _is_str_array(array):
+            indices = cls.items[numpy.isin(cls.names, array)].index
+            return EnumArray(indices, cls)
+
+        # 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 cls.__name__ is array[0].__class__.__name__:
+            indices = cls.items[numpy.isin(cls.enums, array)].index
+            return EnumArray(indices, cls)
+
+        msg = (
+            f"Failed to encode \"{array}\" of type '{array[0].__class__.__name__}', "
+            "as it is not supported. Please, try again with an array of "
+            f"'{int.__name__}', '{str.__name__}', or '{cls.__name__}'."
+        )
+        raise TypeError(msg)
 
 
 __all__ = ["Enum"]

openfisca_core/indexed_enums/enum_array.py

@@ -74,7 +74,7 @@ class EnumArray(t.EnumArray):
 
     def __new__(
         cls,
-        input_array: t.Array[t.DTypeEnum],
+        input_array: t.IndexArray,
         possible_values: None | type[t.Enum] = None,
     ) -> Self:
         """See comment above."""

openfisca_core/indexed_enums/tests/test_enum.py

@@ -27,7 +27,7 @@ def test_enum_encode_with_array_of_enum():
 
 def test_enum_encode_with_enum_sequence():
     """Does encode when called with an enum sequence."""
-    sequence = list(Animal)
+    sequence = list(Animal) + list(Colour)
     enum_array = Animal.encode(sequence)
     assert Animal.DOG in enum_array
 
@@ -89,7 +89,7 @@ def test_enum_encode_with_array_of_string():
 
 def test_enum_encode_with_str_sequence():
     """Does encode when called with a str sequence."""
-    sequence = ("DOG",)
+    sequence = ("DOG", "JAIBA")
     enum_array = Animal.encode(sequence)
     assert Animal.DOG in enum_array
 
@@ -130,5 +130,5 @@ def test_enum_encode_with_any_scalar_array():
 def test_enum_encode_with_any_sequence():
     """Does not encode when called with unsupported types."""
     sequence = memoryview(b"DOG")
-    with pytest.raises(NotImplementedError):
-        Animal.encode(sequence)
+    enum_array = Animal.encode(sequence)
+    assert len(enum_array) == 0