OpenFisca-Core-test 46.0.0__py3-none-any.whl

openfisca_core/indexed_enums/_guards.py

@@ -0,0 +1,209 @@
+from __future__ import annotations
+
+from typing import Final
+
+import numpy
+from typing_extensions import TypeIs
+
+from . import types as t
+
+#: Types for int arrays.
+ints: Final = {
+    numpy.uint8,
+    numpy.uint16,
+    numpy.uint32,
+    numpy.uint64,
+    numpy.int8,
+    numpy.int16,
+    numpy.int32,
+    numpy.int64,
+}
+
+#: Types for object arrays.
+objs: Final = {numpy.object_}
+
+#: Types for str arrays.
+strs: Final = {numpy.str_}
+
+
+def _is_enum_array(array: t.VarArray) -> TypeIs[t.ObjArray]:
+    """Narrow the type of a given array to an array of :obj:`numpy.object_`.
+
+    Args:
+        array: Array to check.
+
+    Returns:
+        bool: True if ``array`` is an array of :obj:`numpy.object_`, False otherwise.
+
+    Examples:
+        >>> import numpy
+
+        >>> from openfisca_core import indexed_enums as enum
+
+        >>> Enum = enum.Enum("Enum", ["A", "B"])
+        >>> array = numpy.array([Enum.A], dtype=numpy.object_)
+        >>> _is_enum_array(array)
+        True
+
+        >>> array = numpy.array([1.0])
+        >>> _is_enum_array(array)
+        False
+
+    """
+    return array.dtype.type in objs
+
+
+def _is_enum_array_like(
+    array: t.VarArray | t.ArrayLike[object],
+) -> TypeIs[t.ArrayLike[t.Enum]]:
+    """Narrow the type of a given array-like to an sequence of :class:`.Enum`.
+
+    Args:
+        array: Array to check.
+
+    Returns:
+        bool: True if ``array`` is an array-like of :class:`.Enum`, False otherwise.
+
+    Examples:
+        >>> from openfisca_core import indexed_enums as enum
+
+        >>> class Housing(enum.Enum):
+        ...     OWNER = "owner"
+        ...     TENANT = "tenant"
+
+        >>> array = [Housing.OWNER]
+        >>> _is_enum_array_like(array)
+        True
+
+        >>> array = ["owner"]
+        >>> _is_enum_array_like(array)
+        False
+
+    """
+    return all(isinstance(item, t.Enum) for item in array)
+
+
+def _is_int_array(array: t.VarArray) -> TypeIs[t.IndexArray]:
+    """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 array.dtype.type in ints
+
+
+def _is_int_array_like(
+    array: t.VarArray | t.ArrayLike[object],
+) -> TypeIs[t.ArrayLike[int]]:
+    """Narrow the type of a given array-like to a sequence of :obj:`int`.
+
+    Args:
+        array: Array to check.
+
+    Returns:
+        bool: True if ``array`` is an array-like of :obj:`int`, False otherwise.
+
+    Examples:
+        >>> array = [1]
+        >>> _is_int_array_like(array)
+        True
+
+        >>> array = (1, 2)
+        >>> _is_int_array_like(array)
+        True
+
+        >>> array = [1.0]
+        >>> _is_int_array_like(array)
+        False
+
+    """
+    return all(isinstance(item, int) for item in array)
+
+
+def _is_str_array(array: t.VarArray) -> 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 array.dtype.type in strs
+
+
+def _is_str_array_like(
+    array: t.VarArray | t.ArrayLike[object],
+) -> TypeIs[t.ArrayLike[str]]:
+    """Narrow the type of a given array-like to an sequence of :obj:`str`.
+
+    Args:
+        array: Array to check.
+
+    Returns:
+        bool: True if ``array`` is an array-like of :obj:`str`, False otherwise.
+
+    Examples:
+        >>> from openfisca_core import indexed_enums as enum
+
+        >>> class Housing(enum.Enum):
+        ...     OWNER = "owner"
+        ...     TENANT = "tenant"
+
+        >>> array = [Housing.OWNER]
+        >>> _is_str_array_like(array)
+        False
+
+        >>> array = ["owner"]
+        >>> _is_str_array_like(array)
+        True
+
+    """
+    return all(isinstance(item, str) for item in array)
+
+
+__all__ = [
+    "_is_enum_array",
+    "_is_enum_array_like",
+    "_is_int_array",
+    "_is_int_array_like",
+    "_is_str_array",
+    "_is_str_array_like",
+]

openfisca_core/indexed_enums/_utils.py

@@ -0,0 +1,180 @@
+from __future__ import annotations
+
+import numpy
+
+from . import types as t
+
+
+def _enum_to_index(value: t.ObjArray | t.ArrayLike[t.Enum]) -> t.IndexArray:
+    """Transform an array of enum members into an index array.
+
+    Args:
+        value: The enum members array to encode.
+
+    Returns:
+        The index array.
+
+    Examples:
+        >>> import numpy
+
+        >>> from openfisca_core import indexed_enums as enum
+
+        >>> class Road(enum.Enum):
+        ...     STREET = (
+        ...         "A public road that connects two points, but also has "
+        ...         "buildings on both sides of it; these typically run "
+        ...         "perpendicular to avenues."
+        ...     )
+        ...     AVENUE = (
+        ...         "A public way that also has buildings and/or trees on both "
+        ...         "sides; these run perpendicular to streets and are "
+        ...         "traditionally wider."
+        ...     )
+
+        >>> class Rogue(enum.Enum):
+        ...     BOULEVARD = "More like a shady impasse, to be honest."
+
+        >>> _enum_to_index(Road.AVENUE)
+        Traceback (most recent call last):
+        TypeError: 'Road' object is not iterable
+
+        >>> _enum_to_index([Road.AVENUE])
+        array([1], dtype=int16)
+
+        >>> _enum_to_index(numpy.array(Road.AVENUE))
+        Traceback (most recent call last):
+        TypeError: iteration over a 0-d array
+
+        >>> _enum_to_index(numpy.array([Road.AVENUE]))
+        array([1], dtype=int16)
+
+        >>> value = numpy.array([Road.STREET, Road.AVENUE, Road.STREET])
+        >>> _enum_to_index(value)
+        array([0, 1, 0], dtype=int16)
+
+        >>> value = numpy.array([Road.AVENUE, Road.AVENUE, Rogue.BOULEVARD])
+        >>> _enum_to_index(value)
+        array([1, 1, 0], dtype=int16)
+
+    """
+    return numpy.array([enum.index for enum in value], t.EnumDType)
+
+
+def _int_to_index(
+    enum_class: type[t.Enum], value: t.IndexArray | t.ArrayLike[int]
+) -> t.IndexArray:
+    """Transform an integer array into an index array.
+
+    Args:
+        enum_class: The enum class to encode the integer array.
+        value: The integer array to encode.
+
+    Returns:
+        The index array.
+
+    Examples:
+        >>> from array import array
+
+        >>> import numpy
+
+        >>> from openfisca_core import indexed_enums as enum
+
+        >>> class Road(enum.Enum):
+        ...     STREET = (
+        ...         "A public road that connects two points, but also has "
+        ...         "buildings on both sides of it; these typically run "
+        ...         "perpendicular to avenues."
+        ...     )
+        ...     AVENUE = (
+        ...         "A public way that also has buildings and/or trees on both "
+        ...         "sides; these run perpendicular to streets and are "
+        ...         "traditionally wider."
+        ...     )
+
+        >>> _int_to_index(Road, 1)
+        array([1], dtype=int16)
+
+        >>> _int_to_index(Road, [1])
+        array([1], dtype=int16)
+
+        >>> _int_to_index(Road, array("B", [1]))
+        array([1], dtype=int16)
+
+        >>> _int_to_index(Road, memoryview(array("B", [1])))
+        array([1], dtype=int16)
+
+        >>> _int_to_index(Road, numpy.array(1))
+        array([1], dtype=int16)
+
+        >>> _int_to_index(Road, numpy.array([1]))
+        array([1], dtype=int16)
+
+        >>> _int_to_index(Road, numpy.array([0, 1, 0]))
+        array([0, 1, 0], dtype=int16)
+
+        >>> _int_to_index(Road, numpy.array([1, 1, 2]))
+        array([1, 1], dtype=int16)
+
+    """
+    indices = enum_class.indices
+    values = numpy.asarray(value)
+    return values[values < indices.size].astype(t.EnumDType)
+
+
+def _str_to_index(
+    enum_class: type[t.Enum], value: t.StrArray | t.ArrayLike[str]
+) -> t.IndexArray:
+    """Transform a string array into an index array.
+
+    Args:
+        enum_class: The enum class to encode the string array.
+        value: The string array to encode.
+
+    Returns:
+        The index array.
+
+    Examples:
+        >>> import numpy
+
+        >>> from openfisca_core import indexed_enums as enum
+
+        >>> class Road(enum.Enum):
+        ...     STREET = (
+        ...         "A public road that connects two points, but also has "
+        ...         "buildings on both sides of it; these typically run "
+        ...         "perpendicular to avenues."
+        ...     )
+        ...     AVENUE = (
+        ...         "A public way that also has buildings and/or trees on both "
+        ...         "sides; these run perpendicular to streets and are "
+        ...         "traditionally wider."
+        ...     )
+
+        >>> _str_to_index(Road, "AVENUE")
+        array([1], dtype=int16)
+
+        >>> _str_to_index(Road, ["AVENUE"])
+        array([1], dtype=int16)
+
+        >>> _str_to_index(Road, numpy.array("AVENUE"))
+        array([1], dtype=int16)
+
+        >>> _str_to_index(Road, numpy.array(["AVENUE"]))
+        array([1], dtype=int16)
+
+        >>> _str_to_index(Road, numpy.array(["STREET", "AVENUE", "STREET"]))
+        array([0, 1, 0], dtype=int16)
+
+        >>> _str_to_index(Road, numpy.array(["AVENUE", "AVENUE", "BOULEVARD"]))
+        array([1, 1], dtype=int16)
+
+    """
+    values = numpy.asarray(value)
+    names = enum_class.names
+    mask = numpy.isin(values, names)
+    sorter = numpy.argsort(names)
+    result = sorter[numpy.searchsorted(names, values[mask], sorter=sorter)]
+    return result.astype(t.EnumDType)
+
+
+__all__ = ["_enum_to_index", "_int_to_index", "_str_to_index"]

openfisca_core/indexed_enums/config.py

@@ -0,0 +1,6 @@
+import numpy
+
+ENUM_ARRAY_DTYPE = numpy.int16
+
+
+__all__ = ["ENUM_ARRAY_DTYPE"]

openfisca_core/indexed_enums/enum.py

@@ -0,0 +1,218 @@
+from __future__ import annotations
+
+from collections.abc import Sequence
+
+import numpy
+
+from . import types as t
+from ._enum_type import EnumType
+from ._errors import EnumEncodingError, EnumMemberNotFoundError
+from ._guards import (
+    _is_enum_array,
+    _is_enum_array_like,
+    _is_int_array,
+    _is_int_array_like,
+    _is_str_array,
+    _is_str_array_like,
+)
+from ._utils import _enum_to_index, _int_to_index, _str_to_index
+from .enum_array import EnumArray
+
+
+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
+    :mod:`~openfisca_core.simulations` on large :mod:`~openfisca_core.populations`.
+
+    Examples:
+        >>> from openfisca_core import indexed_enums as enum
+
+        >>> class Housing(enum.Enum):
+        ...     OWNER = "Owner"
+        ...     TENANT = "Tenant"
+        ...     FREE_LODGER = "Free lodger"
+        ...     HOMELESS = "Homeless"
+
+        >>> repr(Housing)
+        "<enum 'Housing'>"
+
+        >>> repr(Housing.TENANT)
+        'Housing.TENANT'
+
+        >>> str(Housing.TENANT)
+        'Housing.TENANT'
+
+        >>> dict([(Housing.TENANT, Housing.TENANT.value)])
+        {Housing.TENANT: 'Tenant'}
+
+        >>> list(Housing)
+        [Housing.OWNER, Housing.TENANT, Housing.FREE_LODGER, Housing.HOMELESS]
+
+        >>> Housing["TENANT"]
+        Housing.TENANT
+
+        >>> Housing("Tenant")
+        Housing.TENANT
+
+        >>> Housing.TENANT in Housing
+        True
+
+        >>> len(Housing)
+        4
+
+        >>> Housing.TENANT == Housing.TENANT
+        True
+
+        >>> Housing.TENANT != Housing.TENANT
+        False
+
+        >>> Housing.TENANT.index
+        1
+
+        >>> Housing.TENANT.name
+        'TENANT'
+
+        >>> Housing.TENANT.value
+        'Tenant'
+
+    """
+
+    #: The :attr:`index` of the :class:`.Enum` member.
+    index: int
+
+    def __init__(self, *__args: object, **__kwargs: object) -> None:
+        """Tweak :class:`enum.Enum` to add an :attr:`.index` to each enum item.
+
+        When the enum is initialised, ``_member_names_`` contains the names of
+        the already initialized items, so its length is the index of this item.
+
+        Args:
+            *__args: Positional arguments.
+            **__kwargs: Keyword arguments.
+
+        Examples:
+            >>> import numpy
+
+            >>> from openfisca_core import indexed_enums as enum
+
+            >>> Housing = enum.Enum("Housing", "owner tenant")
+            >>> Housing.tenant.index
+            1
+
+            >>> class Housing(enum.Enum):
+            ...     OWNER = "Owner"
+            ...     TENANT = "Tenant"
+
+            >>> Housing.TENANT.index
+            1
+
+            >>> array = numpy.array([[1, 2], [3, 4]])
+            >>> array[Housing.TENANT.index]
+            array([3, 4])
+
+        Note:
+            ``_member_names_`` is undocumented in upstream :class:`enum.Enum`.
+
+        """
+        self.index = len(self._member_names_)
+
+    def __repr__(self) -> str:
+        return f"{self.__class__.__name__}.{self.name}"
+
+    @classmethod
+    def encode(cls, array: t.VarArray | t.ArrayLike[object]) -> t.EnumArray:
+        """Encode an encodable array into an :class:`.EnumArray`.
+
+        Args:
+            array: :class:`~numpy.ndarray` to encode.
+
+        Returns:
+            EnumArray: An :class:`.EnumArray` with the encoded input values.
+
+        Examples:
+            >>> import numpy
+
+            >>> from openfisca_core import indexed_enums as enum
+
+            >>> class Housing(enum.Enum):
+            ...     OWNER = "Owner"
+            ...     TENANT = "Tenant"
+
+            # EnumArray
+
+            >>> array = numpy.array([1])
+            >>> enum_array = enum.EnumArray(array, Housing)
+            >>> Housing.encode(enum_array)
+            EnumArray([Housing.TENANT])
+
+            # Array of Enum
+
+            >>> array = numpy.array([Housing.TENANT])
+            >>> enum_array = Housing.encode(array)
+            >>> enum_array == Housing.TENANT
+            array([ True])
+
+            # Array of integers
+
+            >>> array = numpy.array([1])
+            >>> enum_array = Housing.encode(array)
+            >>> enum_array == Housing.TENANT
+            array([ True])
+
+            # Array of strings
+
+            >>> array = numpy.array(["TENANT"])
+            >>> enum_array = Housing.encode(array)
+            >>> enum_array == Housing.TENANT
+            array([ True])
+
+            # Array of bytes
+
+            >>> array = numpy.array([b"TENANT"])
+            >>> enum_array = Housing.encode(array)
+            Traceback (most recent call last):
+            EnumEncodingError: Failed to encode "[b'TENANT']" of type 'bytes...
+
+        .. seealso::
+            :meth:`.EnumArray.decode` for decoding.
+
+        """
+        if isinstance(array, EnumArray):
+            return array
+        if len(array) == 0:
+            return EnumArray(numpy.asarray(array, t.EnumDType), cls)
+        if isinstance(array, Sequence):
+            return cls._encode_array_like(array)
+        return cls._encode_array(array)
+
+    @classmethod
+    def _encode_array(cls, value: t.VarArray) -> t.EnumArray:
+        if _is_int_array(value):
+            indices = _int_to_index(cls, value)
+        elif _is_str_array(value):  # type: ignore[unreachable]
+            indices = _str_to_index(cls, value)
+        elif _is_enum_array(value) and cls == value[0].__class__:
+            indices = _enum_to_index(value)
+        else:
+            raise EnumEncodingError(cls, value)
+        if indices.size != len(value):
+            raise EnumMemberNotFoundError(cls)
+        return EnumArray(indices, cls)
+
+    @classmethod
+    def _encode_array_like(cls, value: t.ArrayLike[object]) -> t.EnumArray:
+        if _is_int_array_like(value):
+            indices = _int_to_index(cls, value)
+        elif _is_str_array_like(value):  # type: ignore[unreachable]
+            indices = _str_to_index(cls, value)
+        elif _is_enum_array_like(value):
+            indices = _enum_to_index(value)
+        else:
+            raise EnumEncodingError(cls, value)
+        if indices.size != len(value):
+            raise EnumMemberNotFoundError(cls)
+        return EnumArray(indices, cls)
+
+
+__all__ = ["Enum"]