Refactor formula starting date

openfisca_core/indexed_enums/__init__.py

@@ -21,6 +21,8 @@
 #
 # See: https://www.python.org/dev/peps/pep-0008/#imports
 
-from .config import ENUM_ARRAY_DTYPE  # noqa: F401
-from .enum_array import EnumArray  # noqa: F401
-from .enum import Enum  # noqa: F401
+from .config import ENUM_ARRAY_DTYPE
+from .enum_array import EnumArray
+from .enum import Enum
+
+__all__ = ["ENUM_ARRAY_DTYPE", "EnumArray", "Enum"]

openfisca_core/variables/config.py

@@ -2,9 +2,7 @@
 
 import numpy
 
-from openfisca_core import indexed_enums
-from openfisca_core.indexed_enums import Enum
-
+from openfisca_core import indexed_enums as enums
 
 VALUE_TYPES = {
     bool: {
@@ -35,8 +33,8 @@
         "formatted_value_type": "String",
         "is_period_size_independent": True,
     },
-    Enum: {
-        "dtype": indexed_enums.ENUM_ARRAY_DTYPE,
+    enums.Enum: {
+        "dtype": enums.ENUM_ARRAY_DTYPE,
         "json_type": "string",
         "formatted_value_type": "String",
         "is_period_size_independent": True,
@@ -51,4 +49,7 @@
 }
 
 
-FORMULA_NAME_PREFIX = "formula"
+FORMULA_NAME_PREFIX: str = "formula"
+
+# YYYY or YYYY_MM or YYYY_MM_DD
+FORMULA_NAME_REGEX: str = r"formula_(\d{4})(?:_(\d{2}))?(?:_(\d{2}))?$"

openfisca_core/variables/errors.py

@@ -0,0 +1,14 @@
+import os
+
+
+class FormulaNameError(ValueError):
+    """Raised when a formula name is invalid."""
+
+    def __init__(self, name: str):
+        message = (
+            "Unrecognized formula name. Expecting 'formula_YYYY'",
+            "'formula_YYYY_MM', or 'formula_YYYY_MM_DD', where YYYY, MM and",
+            f"DD are year, month and day. Found: '{name}'.",
+        )
+
+        super().__init__(os.linesep.join(message))

openfisca_core/variables/variable.py

@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+from openfisca_core.types import Formula, Instant
 from typing import Optional, Union
 
 import datetime
@@ -13,9 +14,9 @@
 from openfisca_core.entities import Entity
 from openfisca_core.indexed_enums import Enum, EnumArray
 from openfisca_core.periods import DateUnit, Period
-from openfisca_core.types import Formula, Instant
 
 from . import config, helpers
+from .errors import FormulaNameError
 
 
 class Variable:
@@ -338,21 +339,10 @@ def raise_error():
                 )
             )
 
-        if attribute_name == config.FORMULA_NAME_PREFIX:
-            return datetime.date.min
-
-        FORMULA_REGEX = r"formula_(\d{4})(?:_(\d{2}))?(?:_(\d{2}))?$"  # YYYY or YYYY_MM or YYYY_MM_DD
-
-        match = re.match(FORMULA_REGEX, attribute_name)
-        if not match:
-            raise_error()
-        date_str = "-".join(
-            [match.group(1), match.group(2) or "01", match.group(3) or "01"]
-        )
-
         try:
-            return datetime.datetime.strptime(date_str, "%Y-%m-%d").date()
-        except ValueError:  # formula_2005_99_99 for instance
+            return _FormulaStartingDate.from_name(attribute_name)
+
+        except FormulaNameError:
             raise_error()
 
     # ----- Methods ----- #
@@ -473,3 +463,51 @@ def default_array(self, array_size):
             return EnumArray(array, self.possible_values)
         array.fill(self.default_value)
         return array
+
+
+class _FormulaStartingDate(datetime.date):
+    #: Prefix of the formula's name.
+    prefix: str = config.FORMULA_NAME_PREFIX
+
+    #: Regex used to parse the formula's name.
+    regex: str = config.FORMULA_NAME_REGEX
+
+    @classmethod
+    def from_name(cls, name: str) -> _FormulaStartingDate:
+        """Return the starting date of a formula based on its name.
+
+        Valid dated name formats are : "formula", "formula_YYYY",
+        "formula_YYYY_MM" and "formula_YYYY_MM_DD", where YYYY, MM and DD are
+        year, month and day.
+
+        By convention, the starting date of:
+            - `formula` is `0001-01-01` (minimal date in Python)
+            - `formula_YYYY` is `YYYY-01-01`
+            - `formula_YYYY_MM` is `YYYY-MM-01`
+
+        """
+
+        year: int
+        month: int
+        day: int
+        match: re.Match[str] | None
+
+        if name == cls.prefix:
+            year, month, day, *_ = datetime.date.min.timetuple()
+            return cls.__new__(cls, year, month, day)
+
+        match = re.match(cls.regex, name)
+
+        if match is None:
+            raise FormulaNameError(name)
+
+        else:
+            year = int(match.group(1))
+            month = int(match.group(2) or 1)
+            day = int(match.group(3) or 1)
+
+        try:
+            return cls.__new__(cls, year, month, day)
+
+        except ValueError:  # formula_2005_99_99 for instance
+            raise FormulaNameError(name)