diff --git a/openfisca_core/indexed_enums/enum.py b/openfisca_core/indexed_enums/enum.py index 0dba667251..816afc1b33 100644 --- a/openfisca_core/indexed_enums/enum.py +++ b/openfisca_core/indexed_enums/enum.py @@ -1,23 +1,16 @@ from __future__ import annotations import enum -from typing import Union +from typing import Any, TypeVar, Union import numpy -from openfisca_core.types import ArrayType, Choosable +from openfisca_core.types import ArrayType, Encodable from .. import indexed_enums as enums from .enum_array import EnumArray -#: Type of any encodable array. -Encodable = Union[ - EnumArray, - ArrayType[Choosable], - ArrayType[bytes], - ArrayType[int], - ArrayType[str], - ] +T = TypeVar("T", Encodable, bytes, int, str) class Enum(enum.Enum): @@ -26,39 +19,78 @@ class Enum(enum.Enum): Whose items have an :obj:`int` index. This is useful and performant when running simulations on large populations. + Attributes: + index (:obj:`int`): The ``index`` of the :class:`.Enum` member. + name (:obj:`str`): The ``name`` of the :class:`.Enum` member. + value: The ``value`` of the :class:`.Enum` member. + Examples: >>> class Housing(Enum): - ... owner = "Owner" - ... tenant = "Tenant" - ... free_lodger = "Free lodger" - ... homeless = "Homeless" + ... OWNER = "Owner" + ... TENANT = "Tenant" + ... FREE_LODGER = "Free lodger" + ... HOMELESS = "Homeless" + + >>> repr(Housing) + "" + + >>> repr(Housing.TENANT) + '' + + >>> str(Housing.TENANT) + 'Housing.TENANT' + + >>> dict([(Housing.TENANT, Housing.TENANT.value)]) + {: 'Tenant'} + + >>> tuple(Housing) + (, , ...) - >>> Housing - + >>> Housing["TENANT"] + - >>> list(Housing) - [, ...] + >>> Housing("Tenant") + + + >>> Housing.TENANT in Housing + True >>> len(Housing) 4 - >>> Housing.tenant - + >>> Housing.TENANT == Housing.TENANT + True + + >>> Housing.TENANT != Housing.TENANT + False + + >>> Housing.TENANT > Housing.TENANT + False - >>> Housing["tenant"] - + >>> Housing.TENANT < Housing.TENANT + False - >>> Housing.tenant.index + >>> Housing.TENANT >= Housing.TENANT + True + + >>> Housing.TENANT <= Housing.TENANT + True + + >>> Housing.TENANT.index 1 - >>> Housing.tenant.name - 'tenant' + >>> Housing.TENANT.name + 'TENANT' - >>> Housing.tenant.value + >>> Housing.TENANT.value 'Tenant' """ + index: int + name: str + value: Any + def __init__(self, name: str) -> None: """ Tweaks :class:`~enum.Enum` to add an index to each enum item. @@ -81,18 +113,52 @@ def __init__(self, name: str) -> None: >>> MyEnum.bar.index 1 + >>> array = numpy.array([[1, 2], [3, 4]]) + >>> array[MyEnum.bar.index] + array([3, 4]) + """ self.index = len(self._member_names_) - #: Bypass the slow :meth:`~enum.Enum.__eq__`. + def __repr__(self) -> str: + return f"<{self.__class__.__name__}.{self.name}({self.value})>" + + def __str__(self) -> str: + return f"{self.__class__.__name__}.{self.name}" + + def __lt__(self, other: object) -> bool: + if not isinstance(other, Enum): + return NotImplemented + + return self.index < other.index + + def __le__(self, other: object) -> bool: + if not isinstance(other, Enum): + return NotImplemented + + return self.index <= other.index + + def __gt__(self, other: object) -> bool: + if not isinstance(other, Enum): + return NotImplemented + + return self.index > other.index + + def __ge__(self, other: object) -> bool: + if not isinstance(other, Enum): + return NotImplemented + + return self.index >= other.index + __eq__ = object.__eq__ + """Bypass the slow :meth:`~enum.Enum.__eq__`.""" - #: :meth:`.__hash__` must also be defined as so to stay hashable. __hash__ = object.__hash__ + """:meth:`.__hash__` must also be defined as so to stay hashable.""" @classmethod - def encode(cls, array: Encodable) -> EnumArray: + def encode(cls, array: Union[EnumArray, ArrayType[T]]) -> EnumArray: """Encodes an encodable array into an :obj:`.EnumArray`. Args: @@ -112,7 +178,7 @@ def encode(cls, array: Encodable) -> EnumArray: >>> array = numpy.array([1]) >>> enum_array = EnumArray(array, MyEnum) >>> MyEnum.encode(enum_array) - EnumArray([]) + ])> # ArrayTipe[Enum] diff --git a/openfisca_core/indexed_enums/enum_array.py b/openfisca_core/indexed_enums/enum_array.py index 4109db10f3..f4c8f7bfc7 100644 --- a/openfisca_core/indexed_enums/enum_array.py +++ b/openfisca_core/indexed_enums/enum_array.py @@ -4,7 +4,7 @@ import numpy -from openfisca_core.types import ArrayLike, ArrayType, Choosable +from openfisca_core.types import ArrayLike, ArrayType, Encodable class EnumArray(numpy.ndarray): @@ -23,21 +23,46 @@ class EnumArray(numpy.ndarray): >>> from openfisca_core.variables import Variable >>> class Housing(Enum): - ... owner = "Owner" - ... tenant = "Tenant" - ... free_lodger = "Free lodger" - ... homeless = "Homeless" + ... Owner = "Owner" + ... Tenant = "Tenant" + ... FreeLodger = "Free lodger" + ... Homeless = "Homeless" >>> array = numpy.array([1]) - >>> EnumArray(array, Housing) - EnumArray([]) + >>> enum_array = EnumArray(array, Housing) + + >>> repr(EnumArray) + "" + + >>> repr(enum_array) + '])>' + + >>> str(enum_array) + "['Tenant']" + + >>> tuple(enum_array) + (1,) + + >>> enum_array[0] + 1 + + >>> enum_array[0] in enum_array + True + + >>> len(enum_array) + 1 + + >>> enum_array = EnumArray(list(Housing), Housing) + >>> enum_array[Housing.Tenant.index] + >>> class OccupancyStatus(Variable): ... value_type = Enum ... possible_values = Housing >>> EnumArray(array, OccupancyStatus.possible_values) - EnumArray([]) + ])> + .. _Subclassing ndarray: https://numpy.org/doc/stable/user/basics.subclassing.html @@ -47,22 +72,26 @@ class EnumArray(numpy.ndarray): def __new__( cls, input_array: ArrayType[int], - possible_values: Optional[Type[Choosable]] = None, + possible_values: Optional[Type[Encodable]] = None, ) -> EnumArray: - """See comment above…""" + """See comment above.""" obj = numpy.asarray(input_array).view(cls) obj.possible_values = possible_values return obj def __array_finalize__(self, obj: Optional[ArrayType[int]]) -> None: - """See comment above…""" - if obj is None: return self.possible_values = getattr(obj, "possible_values", None) + def __repr__(self) -> str: + return f"<{self.__class__.__name__}({str(self.decode())})>" + + def __str__(self) -> str: + return str(self.decode_to_str()) + def __eq__(self, other: Any) -> Union[ArrayType[bool], bool]: """Compare equality with the item index. @@ -163,7 +192,7 @@ def _forbidden_operation(self, other: Any) -> NoReturn: __and__ = _forbidden_operation __or__ = _forbidden_operation - def decode(self) -> ArrayLike[Choosable]: + def decode(self) -> ArrayLike[Encodable]: """Decodes itself to a normal array. Returns: @@ -179,7 +208,7 @@ def decode(self) -> ArrayLike[Choosable]: >>> array = numpy.array([1]) >>> enum_array = EnumArray(array, MyEnum) >>> enum_array.decode() - array([]...) + array([], dtype=object) """ @@ -212,9 +241,3 @@ def decode_to_str(self) -> ArrayType[str]: [self == item.index for item in self.possible_values], [item.name for item in self.possible_values], ) - - def __repr__(self) -> str: - return f"{self.__class__.__name__}({str(self.decode())})" - - def __str__(self) -> str: - return str(self.decode_to_str()) diff --git a/openfisca_core/indexed_enums/tests/__init__.py b/openfisca_core/indexed_enums/tests/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/openfisca_core/indexed_enums/tests/test_enum_array.py b/openfisca_core/indexed_enums/tests/test_enum_array.py index 5459411690..2a56452a70 100644 --- a/openfisca_core/indexed_enums/tests/test_enum_array.py +++ b/openfisca_core/indexed_enums/tests/test_enum_array.py @@ -35,15 +35,3 @@ def test_enum_array_any_other_operation(enum_array): with pytest.raises(TypeError, match = "Forbidden operation."): enum_array * 1 - - -def test_enum_array___repr__(enum_array): - """Enum arrays have a custom debugging representation.""" - - assert repr(enum_array) == "EnumArray([])" - - -def test_enum_array___str__(enum_array): - """Enum arrays have a custom end-user representation.""" - - assert str(enum_array) == "['bar']" diff --git a/openfisca_core/types/__init__.py b/openfisca_core/types/__init__.py index 4216f13924..0d2da9673b 100644 --- a/openfisca_core/types/__init__.py +++ b/openfisca_core/types/__init__.py @@ -34,10 +34,11 @@ from .protocols import ( # noqa: F401 Buildable, - Choosable, Descriptable, + Encodable, Modelable, Personifiable, Representable, Rolifiable, + Timeable, ) diff --git a/openfisca_core/types/callables/formulas.py b/openfisca_core/types/callables/formulas.py index ce7720072d..2e60a97627 100644 --- a/openfisca_core/types/callables/formulas.py +++ b/openfisca_core/types/callables/formulas.py @@ -1,8 +1,8 @@ from typing import Callable +from ..protocols.timeable import Timeable from ..protocols._aggregatable import Aggregatable from ..protocols._instantizable import Instantizable -from ..protocols._timeable import Timeable from ..data_types import ArrayType ParamsType = Callable[[Timeable], Instantizable] diff --git a/openfisca_core/types/protocols/__init__.py b/openfisca_core/types/protocols/__init__.py index a4c1a51ce3..73893e5719 100644 --- a/openfisca_core/types/protocols/__init__.py +++ b/openfisca_core/types/protocols/__init__.py @@ -22,9 +22,10 @@ # See: https://www.python.org/dev/peps/pep-0008/#imports from .buildable import Buildable # noqa: F401 -from .choosable import Choosable # noqa: F401 from .descriptable import Descriptable # noqa: F401 +from .encodable import Encodable # noqa: F401 from .modelable import Modelable # noqa: F401 from .personifiable import Personifiable # noqa: F401 from .representable import Representable # noqa: F401 from .rolifiable import Rolifiable # noqa: F401 +from .timeable import Timeable # noqa: F401 diff --git a/openfisca_core/types/protocols/buildable.py b/openfisca_core/types/protocols/buildable.py index 84e11f0511..c61b525e6b 100644 --- a/openfisca_core/types/protocols/buildable.py +++ b/openfisca_core/types/protocols/buildable.py @@ -28,9 +28,11 @@ def __init__(self, builder: RT, buildee: Type[ET]) -> None: @abc.abstractmethod def __call__(self, items: Iterable[EL]) -> Sequence[ET]: """A concrete builder implements :meth:`.__call__`.""" + ... @abc.abstractmethod def build(self, item: EL) -> ET: """A concrete builder implements :meth:`.build`.""" + ... diff --git a/openfisca_core/types/protocols/encodable.py b/openfisca_core/types/protocols/encodable.py new file mode 100644 index 0000000000..ce33c1cd6a --- /dev/null +++ b/openfisca_core/types/protocols/encodable.py @@ -0,0 +1,23 @@ +import abc +from typing import Any + +from typing_extensions import Protocol + + +class Encodable(Protocol): + """Base type for any model implementing a literal list of choices. + + Type-checking against abstractions rather than implementations helps in + (a) decoupling the codebse, thanks to structural subtyping, and + (b) documenting/enforcing the blueprints of the different OpenFisca models. + + .. versionadded:: 35.8.0 + + """ + + @classmethod + @abc.abstractmethod + def encode(cls, array: Any) -> Any: + """A concrete encodable model implements :meth:`.encode`.""" + + ... diff --git a/openfisca_core/types/protocols/representable.py b/openfisca_core/types/protocols/representable.py index ff3a26b48f..48557b3d7f 100644 --- a/openfisca_core/types/protocols/representable.py +++ b/openfisca_core/types/protocols/representable.py @@ -18,4 +18,5 @@ class Representable(Protocol): @abc.abstractmethod def get_variable(self, name: str, check: bool) -> Optional[Any]: """A concrete representable model implements :meth:`.get_variable`.""" + ... diff --git a/openfisca_core/types/protocols/timeable.py b/openfisca_core/types/protocols/timeable.py new file mode 100644 index 0000000000..0303955eec --- /dev/null +++ b/openfisca_core/types/protocols/timeable.py @@ -0,0 +1,15 @@ +from typing_extensions import Protocol + + +class Timeable(Protocol): + """Base type for any model implementing a period/instant-like behaviour. + + Type-checking against abstractions rather than implementations helps in + (a) decoupling the codebse, thanks to structural subtyping, and + (b) documenting/enforcing the blueprints of the different OpenFisca models. + + .. versionadded:: 35.7.0 + + """ + + ... diff --git a/setup.cfg b/setup.cfg index ed5d56349b..e8768597d9 100644 --- a/setup.cfg +++ b/setup.cfg @@ -35,5 +35,8 @@ ignore_errors = True [mypy-openfisca_core.entities.tests.*] ignore_errors = True +[mypy-openfisca_core.indexed_enums.tests.*] +ignore_errors = True + [mypy-openfisca_core.scripts.*] ignore_errors = True diff --git a/tests/core/test_tracers.py b/tests/core/test_tracers.py index 2e3d8dbb56..ef08992b3c 100644 --- a/tests/core/test_tracers.py +++ b/tests/core/test_tracers.py @@ -409,7 +409,7 @@ def test_log_aggregate_with_enum(tracer): tracer._exit_calculation() lines = tracer.computation_log.lines(aggregate = True) - assert lines[0] == " A<2017> >> {'avg': EnumArray(HousingOccupancyStatus.tenant), 'max': EnumArray(HousingOccupancyStatus.tenant), 'min': EnumArray(HousingOccupancyStatus.tenant)}" + assert lines[0] == " A<2017> >> {'avg': , 'max': , 'min': }" def test_log_aggregate_with_strings(tracer):