refactor: consolidate commons module types

openfisca_core/commons/__init__.py

@@ -1,6 +1,5 @@
 """Common tools for contributors and users."""
 
-from . import types
 from .dummy import Dummy
 from .formulas import apply_thresholds, concat, switch
 from .misc import empty_clone, eval_expression, stringify_array
@@ -16,5 +15,4 @@
     "marginal_rate",
     "stringify_array",
     "switch",
-    "types",
 ]

openfisca_core/commons/formulas.py

@@ -4,14 +4,14 @@
 
 import numpy
 
-from . import types as t
+from openfisca_core import types as t
 
 
 def apply_thresholds(
-    input: t.Array[numpy.float32],
+    input: t.FloatArray,
     thresholds: t.ArrayLike[float],
     choices: t.ArrayLike[float],
-) -> t.Array[numpy.float32]:
+) -> t.FloatArray:
     """Makes a choice based on an input and thresholds.
 
     From a list of ``choices``, this function selects one of these values
@@ -52,9 +52,9 @@ def apply_thresholds(
 
 
 def concat(
-    this: t.Array[numpy.str_] | t.ArrayLike[object],
-    that: t.Array[numpy.str_] | t.ArrayLike[object],
-) -> t.Array[numpy.str_]:
+    this: t.StrArray | t.ArrayLike[object],
+    that: t.StrArray | t.ArrayLike[object],
+) -> t.StrArray:
     """Concatenate the values of two arrays.
 
     Args:
@@ -87,9 +87,9 @@ def concat(
 
 
 def switch(
-    conditions: t.Array[numpy.float32] | t.ArrayLike[float],
+    conditions: t.FloatArray | t.ArrayLike[float],
     value_by_condition: Mapping[float, float],
-) -> t.Array[numpy.float32]:
+) -> t.FloatArray:
     """Mimic a switch statement.
 
     Given an array of conditions, returns an array of the same size,

openfisca_core/commons/misc.py

@@ -1,7 +1,6 @@
 from __future__ import annotations
 
 import numexpr
-import numpy
 
 from openfisca_core import types as t
 
@@ -43,7 +42,7 @@ def __init__(_: object) -> None: ...
     return new
 
 
-def stringify_array(array: None | t.Array[numpy.generic]) -> str:
+def stringify_array(array: None | t.VarArray) -> str:
     """Generate a clean string representation of a numpy array.
 
     Args:
@@ -79,7 +78,7 @@ def stringify_array(array: None | t.Array[numpy.generic]) -> str:
 
 def eval_expression(
     expression: str,
-) -> str | t.Array[numpy.bool_] | t.Array[numpy.int32] | t.Array[numpy.float32]:
+) -> str | t.BoolArray | t.IntArray | t.FloatArray:
     """Evaluate a string expression to a numpy array.
 
     Args:

openfisca_core/commons/rates.py

@@ -2,14 +2,14 @@
 
 import numpy
 
-from . import types as t
+from openfisca_core import types as t
 
 
 def average_rate(
-    target: t.Array[numpy.float32],
-    varying: t.Array[numpy.float32] | t.ArrayLike[float],
+    target: t.FloatArray,
+    varying: t.FloatArray | t.ArrayLike[float],
     trim: None | t.ArrayLike[float] = None,
-) -> t.Array[numpy.float32]:
+) -> t.FloatArray:
     """Compute the average rate of a target net income.
 
     Given a ``target`` net income, and according to the ``varying`` gross
@@ -59,10 +59,10 @@ def average_rate(
 
 
 def marginal_rate(
-    target: t.Array[numpy.float32],
-    varying: t.Array[numpy.float32] | t.ArrayLike[float],
+    target: t.FloatArray,
+    varying: t.FloatArray | t.ArrayLike[float],
     trim: None | t.ArrayLike[float] = None,
-) -> t.Array[numpy.float32]:
+) -> t.FloatArray:
     """Compute the marginal rate of a target net income.
 
     Given a ``target`` net income, and according to the ``varying`` gross

openfisca_core/types.py

@@ -10,6 +10,17 @@
 
 import numpy
 import pendulum
+from numpy import (
+    bool_ as BoolDType,
+    bytes_ as BytesDType,
+    datetime64 as DateDType,
+    float32 as FloatDType,
+    generic as VarDType,
+    int32 as IntDType,
+    object_ as ObjDType,
+    str_ as StrDType,
+    uint8 as EnumDType,
+)
 
 #: Generic covariant type var.
 _T_co = TypeVar("_T_co", covariant=True)
@@ -22,6 +33,33 @@
 #: Type representing an numpy array.
 Array: TypeAlias = NDArray[_N_co]
 
+#: Type alias for a boolean array.
+BoolArray: TypeAlias = Array[BoolDType]
+
+#: Type alias for an array of bytes.
+BytesArray: TypeAlias = Array[BytesDType]
+
+#: Type alias for an array of dates.
+DateArray: TypeAlias = Array[DateDType]
+
+#: Type alias for an array of floats.
+FloatArray: TypeAlias = Array[FloatDType]
+
+#: Type alias for an array of enum indices.
+IndexArray: TypeAlias = Array[EnumDType]
+
+#: Type alias for an array of integers.
+IntArray: TypeAlias = Array[IntDType]
+
+#: Type alias for an array of objects.
+ObjArray: TypeAlias = Array[ObjDType]
+
+#: Type alias for an array of strings.
+StrArray: TypeAlias = Array[StrDType]
+
+#: Type alias for an array of generic objects.
+VarArray: TypeAlias = Array[VarDType]
+
 #: Type var for array-like objects.
 _L = TypeVar("_L")