From dbbd16e2f411acb5059aed49021b5e67154138b2 Mon Sep 17 00:00:00 2001 From: Mauko Quiroga Date: Tue, 19 Nov 2024 23:34:54 +0100 Subject: [PATCH] refactor: consolidate commons module types --- openfisca_core/commons/__init__.py | 2 -- openfisca_core/commons/formulas.py | 16 ++++++------- openfisca_core/commons/misc.py | 5 ++-- openfisca_core/commons/rates.py | 14 +++++------ openfisca_core/commons/types.py | 3 --- openfisca_core/types.py | 38 ++++++++++++++++++++++++++++++ 6 files changed, 55 insertions(+), 23 deletions(-) delete mode 100644 openfisca_core/commons/types.py diff --git a/openfisca_core/commons/__init__.py b/openfisca_core/commons/__init__.py index 550088141..a80213430 100644 --- a/openfisca_core/commons/__init__.py +++ b/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", ] diff --git a/openfisca_core/commons/formulas.py b/openfisca_core/commons/formulas.py index 1df803941..bced000a6 100644 --- a/openfisca_core/commons/formulas.py +++ b/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, diff --git a/openfisca_core/commons/misc.py b/openfisca_core/commons/misc.py index e3e55948d..4e0d42ad7 100644 --- a/openfisca_core/commons/misc.py +++ b/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: diff --git a/openfisca_core/commons/rates.py b/openfisca_core/commons/rates.py index 1d17c7735..f2bfe0057 100644 --- a/openfisca_core/commons/rates.py +++ b/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 diff --git a/openfisca_core/commons/types.py b/openfisca_core/commons/types.py deleted file mode 100644 index 39c067f45..000000000 --- a/openfisca_core/commons/types.py +++ /dev/null @@ -1,3 +0,0 @@ -from openfisca_core.types import Array, ArrayLike - -__all__ = ["Array", "ArrayLike"] diff --git a/openfisca_core/types.py b/openfisca_core/types.py index 9c8105741..60573ccd3 100644 --- a/openfisca_core/types.py +++ b/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")