From bc8c16657b945153442b70575c458ca7fc00f7b3 Mon Sep 17 00:00:00 2001
From: Mauko Quiroga <public@mauko.me>
Date: Mon, 14 Oct 2024 17:38:11 +0200
Subject: [PATCH] test: check period (#1238 #1276 #1277)

---
 openfisca_core/populations/__init__.py        |  3 +-
 .../populations/_core_population.py           | 50 +++++++++++++------
 openfisca_core/populations/_errors.py         | 24 ++++++++-
 3 files changed, 60 insertions(+), 17 deletions(-)

diff --git a/openfisca_core/populations/__init__.py b/openfisca_core/populations/__init__.py
index 3fc686a69..e223cd81d 100644
--- a/openfisca_core/populations/__init__.py
+++ b/openfisca_core/populations/__init__.py
@@ -31,7 +31,7 @@
 
 from . import types
 from ._core_population import CorePopulation
-from ._errors import InvalidArraySizeError
+from ._errors import InvalidArraySizeError, PeriodValidityError
 from .config import ADD, DIVIDE
 from .group_population import GroupPopulation
 from .population import Population
@@ -46,6 +46,7 @@
     "FirstPersonToEntityProjector",
     "GroupPopulation",
     "InvalidArraySizeError",
+    "PeriodValidityError",
     "Population",
     "Projector",
     "SinglePopulation",
diff --git a/openfisca_core/populations/_core_population.py b/openfisca_core/populations/_core_population.py
index 708f5f4ef..8dff04137 100644
--- a/openfisca_core/populations/_core_population.py
+++ b/openfisca_core/populations/_core_population.py
@@ -11,6 +11,7 @@
 
 from . import types as t
 from ._errors import InvalidArraySizeError
+from ._errors import PeriodValidityError
 
 #: Type variable for a covariant data type.
 _DT_co = TypeVar("_DT_co", covariant=True, bound=t.VarDType)
@@ -164,26 +165,45 @@ def check_array_compatible_with_entity(self, array: t.FloatArray) -> None:
             return
         raise InvalidArraySizeError(array, self.entity.key, self.count)
 
+    @staticmethod
     def check_period_validity(
-        self,
-        variable_name: str,
-        period: None | t.PeriodLike,
+        variable_name: t.VariableName,
+        period: None | t.PeriodLike = None,
     ) -> None:
+        """Check if a period is valid.
+
+        Args:
+            variable_name: The name of the variable.
+            period: The period to check.
+
+        Raises:
+            PeriodValidityError: If the period is not valid.
+
+        Examples:
+            >>> from openfisca_core import entities, periods, populations
+
+            >>> class Person(entities.SingleEntity): ...
+
+            >>> person = Person("person", "people", "", "")
+            >>> period = periods.Period("2017-04")
+            >>> population = populations.CorePopulation(person)
+
+            >>> population.check_period_validity("salary")
+            Traceback (most recent call last):
+            PeriodValidityError: You requested computation of variable "sala...
+
+            >>> population.check_period_validity("salary", 2017)
+
+            >>> population.check_period_validity("salary", "2017-04")
+
+            >>> population.check_period_validity("salary", period)
+
+        """
         if isinstance(period, (int, str, periods.Period)):
             return
-
         stack = traceback.extract_stack()
-        filename, line_number, function_name, line_of_code = stack[-3]
-        msg = f"""
-You requested computation of variable "{variable_name}", but you did not specify on which period in "{filename}:{line_number}":
-    {line_of_code}
-When you request the computation of a variable within a formula, you must always specify the period as the second parameter. The convention is to call this parameter "period". For example:
-    computed_salary = person('salary', period).
-See more information at <https://openfisca.org/doc/coding-the-legislation/35_periods.html#periods-in-variable-definition>.
-"""
-        raise ValueError(
-            msg,
-        )
+        filename, line_number, _, line_of_code = stack[-3]
+        raise PeriodValidityError(variable_name, filename, line_number, line_of_code)
 
     # Helpers
 
diff --git a/openfisca_core/populations/_errors.py b/openfisca_core/populations/_errors.py
index a664ad71f..77e6c424b 100644
--- a/openfisca_core/populations/_errors.py
+++ b/openfisca_core/populations/_errors.py
@@ -12,4 +12,26 @@ def __init__(self, array: t.FloatArray, entity: t.EntityKey, count: int) -> None
         super().__init__(msg)
 
 
-__all__ = ["InvalidArraySizeError"]
+class PeriodValidityError(ValueError):
+    """Raised when a period is not valid."""
+
+    def __init__(
+        self,
+        variable_name: t.VariableName,
+        filename: str,
+        line_number: int,
+        line_of_code: int,
+    ) -> None:
+        msg = (
+            f'You requested computation of variable "{variable_name}", but '
+            f'you did not specify on which period in "{filename}:{line_number}": '
+            f"{line_of_code}. When you request the computation of a variable "
+            "within a formula, you must always specify the period as the second "
+            'parameter. The convention is to call this parameter "period". For '
+            'example: computed_salary = person("salary", period). More information at '
+            "<https://openfisca.org/doc/coding-the-legislation/35_periods.html#periods-in-variable-definition>."
+        )
+        super().__init__(msg)
+
+
+__all__ = ["InvalidArraySizeError", "PeriodValidityError"]