Fix rates tests

openfisca · Aug 31, 2021 · 8f004a9 · 8f004a9
1 parent 9a6ec1f
commit 8f004a9
Show file tree

Hide file tree

Showing 6 changed files with 21 additions and 15 deletions.
diff --git a/openfisca_core/commons/__init__.py b/openfisca_core/commons/__init__.py
@@ -21,7 +21,7 @@
 #
 # See: https://www.python.org/dev/peps/pep-0008/#imports
 
-from .typing import ArrayLike, NDArray  # noqa: F401
+from .typing import ArrayLike, Array  # noqa: F401
 from .formulas import apply_thresholds, concat, switch  # noqa: F401
 from .misc import empty_clone, stringify_array  # noqa: F401
 from .rates import average_rate, marginal_rate  # noqa: F401

diff --git a/openfisca_core/commons/formulas.py b/openfisca_core/commons/formulas.py
@@ -1,9 +1,9 @@
 import numpy
 
-from .typing import ArrayLike, NDArray
+from .typing import ArrayLike, Array
 
 
-def apply_thresholds(input: NDArray[float], thresholds: ArrayLike[float], choices: ArrayLike[float]) -> NDArray[float]:
+def apply_thresholds(input: Array[float], thresholds: ArrayLike[float], choices: ArrayLike[float]) -> Array[float]:
     """Makes a choice based on an input and thresholds.
 
     From list of ``choices``, it selects one of them based from a list of
@@ -40,7 +40,7 @@ def apply_thresholds(input: NDArray[float], thresholds: ArrayLike[float], choice
     return numpy.select(condlist, choices)
 
 
-def concat(this: ArrayLike[str], that: ArrayLike[str]) -> NDArray[str]:
+def concat(this: ArrayLike[str], that: ArrayLike[str]) -> Array[str]:
     """Concatenates the values of two arrays.
 
     Args:
@@ -66,7 +66,7 @@ def concat(this: ArrayLike[str], that: ArrayLike[str]) -> NDArray[str]:
     return numpy.char.add(this, that)
 
 
-def switch(conditions: NDArray[float], value_by_condition: dict) -> NDArray[float]:
+def switch(conditions: Array[float], value_by_condition: dict) -> Array[float]:
     """Reproduces 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
@@ -1,6 +1,6 @@
 from typing import Any
 
-from .typing import NDArray
+from .typing import Array
 
 
 def empty_clone(original: Any) -> Any:
@@ -40,7 +40,7 @@ def __init__(self) -> None:
     return new
 
 
-def stringify_array(array: NDArray) -> str:
+def stringify_array(array: Array) -> str:
     """Generates a clean string representation of a numpy array.
 
     Args:

diff --git a/openfisca_core/commons/rates.py b/openfisca_core/commons/rates.py
@@ -2,16 +2,19 @@
 
 import numpy
 
-from .typing import ArrayLike, NDArray
+from .typing import ArrayLike, Array
 
 
-def average_rate(target: NDArray[float], varying: ArrayLike[float], trim: Optional[ArrayLike[float]] = None) -> NDArray[float]:
+def average_rate(target: Array[float], varying: ArrayLike[float], trim: Optional[ArrayLike[float]] = None) -> Array[float]:
     """Computes the average rate of a target net income.
 
     Given a ``target`` net income, and according to the ``varying`` gross
     income. Optionally, a ``trim`` can be applied consisting on the lower and
     upper bounds of the average rate to be computed.
 
+    Note:
+        Usually, ``target`` and ``varying`` are the same size.
+
     Args:
         target: The targeted net income.
         varying: The varying gross income.
@@ -25,7 +28,7 @@ def average_rate(target: NDArray[float], varying: ArrayLike[float], trim: Option
 
     Examples:
         >>> target = numpy.array([1, 2, 3])
-        >>> varying = 2
+        >>> varying = [2, 2, 2]
         >>> trim = [-1, .25]
         >>> average_rate(target, varying, trim)
         array([ nan,  0. , -0.5])
@@ -40,13 +43,16 @@ def average_rate(target: NDArray[float], varying: ArrayLike[float], trim: Option
     return average_rate
 
 
-def marginal_rate(target: NDArray[float], varying: NDArray[float], trim: Optional[ArrayLike[float]] = None) -> NDArray[float]:
+def marginal_rate(target: Array[float], varying: Array[float], trim: Optional[ArrayLike[float]] = None) -> Array[float]:
     """Computes the marginal rate of a target net income.
 
     Given a ``target`` net income, and according to the ``varying`` gross
     income. Optionally, a ``trim`` can be applied consisting of the lower and
     upper bounds of the marginal rate to be computed.
 
+    Note:
+        Usually, ``target`` and ``varying`` are the same size.
+
     Args:
         target: The targeted net income.
         varying: The varying gross income.

diff --git a/openfisca_core/commons/tests/test_rates.py b/openfisca_core/commons/tests/test_rates.py
@@ -8,7 +8,7 @@ def test_average_rate_when_varying_is_zero():
     """Yields infinity when the varying gross income crosses zero."""
 
     target = numpy.array([1, 2, 3])
-    varying = 0
+    varying = [0, 0, 0]
 
     result = commons.average_rate(target, varying)
 
@@ -19,7 +19,7 @@ def test_marginal_rate_when_varying_is_zero():
     """Yields infinity when the varying gross income crosses zero."""
 
     target = numpy.array([1, 2, 3])
-    varying = numpy.array([0, 0])
+    varying = numpy.array([0, 0, 0])
 
     result = commons.marginal_rate(target, varying)
 

diff --git a/openfisca_core/commons/typing.py b/openfisca_core/commons/typing.py
@@ -1,12 +1,12 @@
 from typing import Sequence, TypeVar, Union
 
-from nptyping import NDArray
+from nptyping import NDArray as Array
 
 
 _T = TypeVar("_T", float, str)
 """:obj:`.TypeVar`: A wildcard type used by other types."""
 
-ArrayLike = Union[Union[NDArray[float], NDArray[str]], Sequence[_T]]
+ArrayLike = Union[Union[Array[float], Array[str]], Sequence[_T]]
 """:obj:`.Generic`: Type of any castable to :class:`.ndarray`.
 
 These include any :obj:`.ndarray` and sequences (like