test(enums): fix enums doctests

CHANGELOG.md

@@ -1,5 +1,52 @@
 # Changelog
 
+## 42.1.0 [#1273](https://github.com/openfisca/openfisca-core/pull/1273)
+
+#### New features
+
+- Introduce `indexed_enums.EnumType`
+  - Allows for actually fancy indexing `indexed_enums.Enum`
+
+#### Technical changes
+
+- Fix doctests
+  - Now `pytest openfisca_core/indexed_enums` runs without errors
+- Fix bug in `Enum.encode` when passing a scalar
+  - Still raises `TypeError` but with an explanation of why it fails
+- Fix bug in `Enum.encode` when encoding values not present in the enum
+  - When encoding values not present in an enum, `Enum.encode` always encoded 
+    the first item of the enum
+  - Now, it correctly encodes only the values requested that exist in the enum
+
+##### Before
+
+```python
+from openfisca_core import indexed_enums as enum
+
+class TestEnum(enum.Enum):
+     ONE = "one"
+     TWO = "two"
+
+TestEnum.encode([2])
+#  EnumArray([0])
+```
+
+##### After
+
+```python
+from openfisca_core import indexed_enums as enum
+
+class TestEnum(enum.Enum):
+     ONE = "one"
+     TWO = "two"
+
+TestEnum.encode([2])
+#  EnumArray([])
+
+TestEnum.encode([0,1,2,5])
+# EnumArray([<TestEnum.ONE: 'one'> <TestEnum.TWO: 'two'>])
+```
+
 ### 42.0.8 [#1272](https://github.com/openfisca/openfisca-core/pull/1272)
 
 #### Documentation

openfisca_core/data_storage/on_disk_storage.py

@@ -87,7 +87,7 @@ def _decode_file(self, file: str) -> t.Array[t.DTypeGeneric]:
             ...     storage = data_storage.OnDiskStorage(directory)
             ...     storage.put(value, period)
             ...     storage._decode_file(storage._files[period])
-            EnumArray([<Housing.TENANT: 'Tenant'>])
+            EnumArray(Housing.TENANT)
 
         """
         enum = self._enums.get(file)

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,116 @@
+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: ...
+
+        >>> class Housing(Enum):
+        ...     OWNER = "Owner"
+        ...     TENANT = "Tenant"
+
+        >>> Housing.items
+        rec.array([(0, 'OWNER', Housing.OWNER), (1, 'TENANT', Housing.TENAN...)
+
+        >>> Housing.indices
+        array([0, 1], dtype=int16)
+
+        >>> Housing.names
+        array(['OWNER', 'TENANT'], dtype='<U6')
+
+        >>> Housing.enums
+        array([Housing.OWNER, Housing.TENANT], dtype=object)
+
+    """
+
+    #: 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__()})
+
+
+__all__ = ["EnumType"]

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"]