refactor(enums): move config to general (#1233)

openfisca_core/config.py

@@ -0,0 +1,68 @@
+from __future__ import annotations
+
+from typing import Final, final
+
+import dataclasses
+import datetime
+
+from openfisca_core import indexed_enums as enum
+
+from . import types as t
+
+
+@final
+@dataclasses.dataclass(frozen=True)
+class ValueType:
+    dtype: t.DTypeLike
+    default: object
+    json_type: str
+    formatted_value_type: str
+    is_period_size_independent: bool
+
+
+value_types: Final = {
+    bool: ValueType(
+        dtype=t.DTypeBool,
+        default=False,
+        json_type="boolean",
+        formatted_value_type="Boolean",
+        is_period_size_independent=True,
+    ),
+    int: ValueType(
+        dtype=t.DTypeInt,
+        default=0,
+        json_type="integer",
+        formatted_value_type="Int",
+        is_period_size_independent=False,
+    ),
+    float: ValueType(
+        dtype=t.DTypeFloat,
+        default=0,
+        json_type="number",
+        formatted_value_type="Float",
+        is_period_size_independent=False,
+    ),
+    str: ValueType(
+        dtype=t.DTypeStr,
+        default="",
+        json_type="string",
+        formatted_value_type="String",
+        is_period_size_independent=True,
+    ),
+    enum.Enum: ValueType(
+        dtype=t.DTypeEnum,
+        default=None,
+        json_type="string",
+        formatted_value_type="String",
+        is_period_size_independent=True,
+    ),
+    datetime.date: ValueType(
+        dtype="datetime64[D]",
+        default=None,
+        json_type="string",
+        formatted_value_type="Date",
+        is_period_size_independent=True,
+    ),
+}
+
+__all__ = ["value_types"]

openfisca_core/entities/types.py

@@ -33,7 +33,6 @@ class RoleParams(TypedDict, total=False):
     "GroupEntity",
     "Role",
     "RoleKey",
-    "RoleParams",
     "RolePlural",
     "SingleEntity",
     "TaxBenefitSystem",

openfisca_core/indexed_enums/enum.py

@@ -29,10 +29,7 @@ def __init__(self, name: str) -> None:
     def encode(
         cls,
         array: (
-            EnumArray
-            | t.Array[numpy.str_]
-            | t.Array[numpy.int16]
-            | t.Array[numpy.int32]
+            EnumArray | t.Array[t.DTypeEnum] | t.Array[t.DTypeInt] | t.Array[t.DTypeStr]
         ),
     ) -> EnumArray:
         """Encode a string numpy array, an enum item numpy array, or an int numpy
@@ -68,7 +65,7 @@ def encode(
             array = numpy.select(
                 [array == item.name for item in cls],
                 [item.index for item in cls],
-            ).astype(numpy.int16)
+            ).astype(t.DTypeEnum)
 
         # Enum items arrays
         elif array.dtype.kind == "O":
@@ -92,9 +89,9 @@ def encode(
             array = numpy.select(
                 [array == item for item in klass],
                 [item.index for item in klass],
-            ).astype(numpy.int16)
+            ).astype(t.DTypeEnum)
 
-        array = numpy.asarray(array, dtype=numpy.int16)
+        array = numpy.asarray(array, dtype=t.DTypeEnum)
         return EnumArray(array, cls)
 
 

openfisca_core/indexed_enums/enum_array.py

@@ -32,28 +32,28 @@ def __new__(
 
     # See previous comment
     def __array_finalize__(
-        self, obj: None | t.EnumArray | t.Array[numpy.generic]
+        self, obj: None | t.EnumArray | t.Array[t.DTypeObject]
     ) -> None:
         if obj is None:
-            return None
+            return
         if isinstance(obj, EnumArray):
             self.possible_values = obj.possible_values
-        return None
+        return
 
     @overload  # type: ignore[override]
-    def __eq__(self, other: None | t.Enum | type[t.Enum]) -> t.Array[numpy.bool_]: ...
+    def __eq__(self, other: None | t.Enum | type[t.Enum]) -> t.Array[t.DTypeBool]: ...
 
     @overload
-    def __eq__(self, other: object) -> t.Array[numpy.bool_] | bool: ...
+    def __eq__(self, other: object) -> t.Array[t.DTypeBool] | bool: ...
 
-    def __eq__(self, other: object) -> t.Array[numpy.bool_] | bool:
-        boolean_array: t.Array[numpy.bool_]
+    def __eq__(self, other: object) -> t.Array[t.DTypeBool] | bool:
+        boolean_array: t.Array[t.DTypeBool]
         boolean: bool
 
         if self.possible_values is None:
             return NotImplemented
 
-        view: t.Array[numpy.int16] = self.view(numpy.ndarray)
+        view: t.Array[t.DTypeEnum] = self.view(numpy.ndarray)
 
         if other is None or self._is_an_enum_type(other):
             boolean_array = view == other
@@ -71,12 +71,12 @@ def __eq__(self, other: object) -> t.Array[numpy.bool_] | bool:
         return boolean
 
     @overload  # type: ignore[override]
-    def __ne__(self, other: None | t.Enum | type[t.Enum]) -> t.Array[numpy.bool_]: ...
+    def __ne__(self, other: None | t.Enum | type[t.Enum]) -> t.Array[t.DTypeBool]: ...
 
     @overload
-    def __ne__(self, other: object) -> t.Array[numpy.bool_] | bool: ...
+    def __ne__(self, other: object) -> t.Array[t.DTypeBool] | bool: ...
 
-    def __ne__(self, other: object) -> t.Array[numpy.bool_] | bool:
+    def __ne__(self, other: object) -> t.Array[t.DTypeBool] | bool:
         return numpy.logical_not(self == other)
 
     def _forbidden_operation(self, other: object) -> NoReturn:
@@ -97,7 +97,7 @@ def _forbidden_operation(self, other: object) -> NoReturn:
     __and__ = _forbidden_operation  # type: ignore[assignment]
     __or__ = _forbidden_operation  # type: ignore[assignment]
 
-    def decode(self) -> t.Array[numpy.int16]:
+    def decode(self) -> t.Array[t.DTypeEnum]:
         """Return the array of enum items corresponding to self.
 
         For instance:
@@ -116,10 +116,10 @@ def decode(self) -> t.Array[numpy.int16]:
 
         return numpy.select(
             [self == item.index for item in self.possible_values],
-            [item for item in self.possible_values],  # type: ignore[misc]
+            list(self.possible_values),  # pyright: ignore[reportArgumentType]
         )
 
-    def decode_to_str(self) -> t.Array[numpy.str_]:
+    def decode_to_str(self) -> t.Array[t.DTypeStr]:
         """Return the array of string identifiers corresponding to self.
 
         For instance:

openfisca_core/indexed_enums/types.py

@@ -1,3 +1,21 @@
-from openfisca_core.types import Array, DTypeEnum, Enum, EnumArray
+from openfisca_core.types import (
+    Array,
+    DTypeBool,
+    DTypeEnum,
+    DTypeInt,
+    DTypeObject,
+    DTypeStr,
+    Enum,
+    EnumArray,
+)
 
-__all__ = ["Array", "DTypeEnum", "Enum", "EnumArray"]
+__all__ = [
+    "Array",
+    "DTypeBool",
+    "DTypeEnum",
+    "DTypeInt",
+    "DTypeObject",
+    "DTypeStr",
+    "Enum",
+    "EnumArray",
+]

openfisca_core/periods/types.py

@@ -173,11 +173,6 @@ class PeriodStr(str, metaclass=_PeriodStrMeta):  # type: ignore[misc]
 
 __all__ = [
     "DateUnit",
-    "ISOCalendarStr",
-    "ISOFormatStr",
     "Instant",
-    "InstantStr",
     "Period",
-    "PeriodStr",
-    "SeqInt",
 ]

openfisca_core/types.py

@@ -1,7 +1,7 @@
 from __future__ import annotations
 
 from collections.abc import Iterable, Sequence, Sized
-from numpy.typing import NDArray
+from numpy.typing import DTypeLike, NDArray
 from typing import Any, NewType, TypeVar, Union
 from typing_extensions import Protocol, Self, TypeAlias
 
@@ -14,7 +14,7 @@
 #: Generic covariant type var.
 _T_co = TypeVar("_T_co", covariant=True)
 
-# Commons
+# Arrays
 
 #: Type var for numpy arrays.
 _N_co = TypeVar("_N_co", covariant=True, bound="DTypeGeneric")
@@ -105,6 +105,14 @@ def key(self, /) -> RoleKey: ...
     def plural(self, /) -> None | RolePlural: ...
 
 
+# Holders
+
+
+class Holder(Protocol):
+    def clone(self, population: Any, /) -> Holder: ...
+    def get_memory_usage(self, /) -> Any: ...
+
+
 # Indexed enums
 
 
@@ -121,14 +129,6 @@ def __new__(
     ) -> Self: ...
 
 
-# Holders
-
-
-class Holder(Protocol):
-    def clone(self, population: Any, /) -> Holder: ...
-    def get_memory_usage(self, /) -> Any: ...
-
-
 # Parameters
 
 
@@ -239,3 +239,6 @@ def __call__(
 
 class Params(Protocol):
     def __call__(self, instant: Instant, /) -> ParameterNodeAtInstant: ...
+
+
+__all__ = ["DTypeLike"]

openfisca_core/variables/__init__.py

@@ -21,6 +21,6 @@
 #
 # See: https://www.python.org/dev/peps/pep-0008/#imports
 
-from .config import FORMULA_NAME_PREFIX, VALUE_TYPES  # noqa: F401
+from .config import FORMULA_NAME_PREFIX  # noqa: F401
 from .helpers import get_annualized_variable, get_neutralized_variable  # noqa: F401
 from .variable import Variable  # noqa: F401

openfisca_core/variables/config.py

@@ -1,53 +1 @@
-import datetime
-
-import numpy
-
-from openfisca_core import indexed_enums
-from openfisca_core.indexed_enums import Enum
-
-VALUE_TYPES = {
-    bool: {
-        "dtype": numpy.bool_,
-        "default": False,
-        "json_type": "boolean",
-        "formatted_value_type": "Boolean",
-        "is_period_size_independent": True,
-    },
-    int: {
-        "dtype": numpy.int32,
-        "default": 0,
-        "json_type": "integer",
-        "formatted_value_type": "Int",
-        "is_period_size_independent": False,
-    },
-    float: {
-        "dtype": numpy.float32,
-        "default": 0,
-        "json_type": "number",
-        "formatted_value_type": "Float",
-        "is_period_size_independent": False,
-    },
-    str: {
-        "dtype": object,
-        "default": "",
-        "json_type": "string",
-        "formatted_value_type": "String",
-        "is_period_size_independent": True,
-    },
-    Enum: {
-        "dtype": indexed_enums.ENUM_ARRAY_DTYPE,
-        "json_type": "string",
-        "formatted_value_type": "String",
-        "is_period_size_independent": True,
-    },
-    datetime.date: {
-        "dtype": "datetime64[D]",
-        "default": datetime.date.fromtimestamp(0),  # 0 == 1970-01-01
-        "json_type": "string",
-        "formatted_value_type": "Date",
-        "is_period_size_independent": True,
-    },
-}
-
-
 FORMULA_NAME_PREFIX = "formula"