Add date units to periods

openfisca_core/data_storage/in_memory_storage.py

@@ -1,6 +1,7 @@
 import numpy
 
 from openfisca_core import periods
+from openfisca_core.periods import DateUnit
 
 
 class InMemoryStorage:
@@ -14,7 +15,7 @@ def __init__(self, is_eternal = False):
 
     def get(self, period):
         if self.is_eternal:
-            period = periods.period(periods.ETERNITY)
+            period = periods.period(DateUnit.ETERNITY)
         period = periods.period(period)
 
         values = self._arrays.get(period)
@@ -24,7 +25,7 @@ def get(self, period):
 
     def put(self, value, period):
         if self.is_eternal:
-            period = periods.period(periods.ETERNITY)
+            period = periods.period(DateUnit.ETERNITY)
         period = periods.period(period)
 
         self._arrays[period] = value
@@ -35,7 +36,7 @@ def delete(self, period = None):
             return
 
         if self.is_eternal:
-            period = periods.period(periods.ETERNITY)
+            period = periods.period(DateUnit.ETERNITY)
         period = periods.period(period)
 
         self._arrays = {

openfisca_core/data_storage/on_disk_storage.py

@@ -4,6 +4,7 @@
 import numpy
 
 from openfisca_core import periods
+from openfisca_core.periods import DateUnit
 from openfisca_core.indexed_enums import EnumArray
 
 
@@ -28,7 +29,7 @@ def _decode_file(self, file):
 
     def get(self, period):
         if self.is_eternal:
-            period = periods.period(periods.ETERNITY)
+            period = periods.period(DateUnit.ETERNITY)
         period = periods.period(period)
 
         values = self._files.get(period)
@@ -38,7 +39,7 @@ def get(self, period):
 
     def put(self, value, period):
         if self.is_eternal:
-            period = periods.period(periods.ETERNITY)
+            period = periods.period(DateUnit.ETERNITY)
         period = periods.period(period)
 
         filename = str(period)
@@ -55,7 +56,7 @@ def delete(self, period = None):
             return
 
         if self.is_eternal:
-            period = periods.period(periods.ETERNITY)
+            period = periods.period(DateUnit.ETERNITY)
         period = periods.period(period)
 
         if period is not None:

openfisca_core/holders/helpers.py

@@ -2,7 +2,7 @@
 
 import numpy
 
-from openfisca_core import periods
+from openfisca_core.periods import DateUnit
 
 log = logging.getLogger(__name__)
 
@@ -20,10 +20,10 @@ def set_input_dispatch_by_period(holder, period, array):
     period_size = period.size
     period_unit = period.unit
 
-    if holder.variable.definition_period == periods.MONTH:
-        cached_period_unit = periods.MONTH
-    elif holder.variable.definition_period == periods.YEAR:
-        cached_period_unit = periods.YEAR
+    if holder.variable.definition_period == DateUnit.MONTH:
+        cached_period_unit = DateUnit.MONTH
+    elif holder.variable.definition_period == DateUnit.YEAR:
+        cached_period_unit = DateUnit.YEAR
     else:
         raise ValueError('set_input_dispatch_by_period can be used only for yearly or monthly variables.')
 
@@ -55,10 +55,10 @@ def set_input_divide_by_period(holder, period, array):
     period_size = period.size
     period_unit = period.unit
 
-    if holder.variable.definition_period == periods.MONTH:
-        cached_period_unit = periods.MONTH
-    elif holder.variable.definition_period == periods.YEAR:
-        cached_period_unit = periods.YEAR
+    if holder.variable.definition_period == DateUnit.MONTH:
+        cached_period_unit = DateUnit.MONTH
+    elif holder.variable.definition_period == DateUnit.YEAR:
+        cached_period_unit = DateUnit.YEAR
     else:
         raise ValueError('set_input_divide_by_period can be used only for yearly or monthly variables.')
 

openfisca_core/holders/holder.py

@@ -8,6 +8,7 @@
 from openfisca_core.errors import PeriodMismatchError
 from openfisca_core.data_storage import InMemoryStorage, OnDiskStorage
 from openfisca_core.indexed_enums import Enum
+from openfisca_core.periods import DateUnit
 
 
 class Holder:
@@ -19,7 +20,7 @@ def __init__(self, variable, population):
         self.population = population
         self.variable = variable
         self.simulation = population.simulation
-        self._memory_storage = InMemoryStorage(is_eternal = (self.variable.definition_period == periods.ETERNITY))
+        self._memory_storage = InMemoryStorage(is_eternal = (self.variable.definition_period == DateUnit.ETERNITY))
 
         # By default, do not activate on-disk storage, or variable dropping
         self._disk_storage = None
@@ -56,7 +57,7 @@ def create_disk_storage(self, directory = None, preserve = False):
             os.mkdir(storage_dir)
         return OnDiskStorage(
             storage_dir,
-            is_eternal = (self.variable.definition_period == periods.ETERNITY),
+            is_eternal = (self.variable.definition_period == DateUnit.ETERNITY),
             preserve_storage_dir = preserve
             )
 
@@ -148,9 +149,9 @@ def set_input(self, period, array):
         """
 
         period = periods.period(period)
-        if period.unit == periods.ETERNITY and self.variable.definition_period != periods.ETERNITY:
+        if period.unit == DateUnit.ETERNITY and self.variable.definition_period != DateUnit.ETERNITY:
             error_message = os.linesep.join([
-                'Unable to set a value for variable {0} for periods.ETERNITY.',
+                'Unable to set a value for variable {0} for DateUnit.ETERNITY.',
                 '{0} is only defined for {1}s. Please adapt your input.',
                 ]).format(
                     self.variable.name,
@@ -197,9 +198,9 @@ def _to_array(self, value):
 
     def _set(self, period, value):
         value = self._to_array(value)
-        if self.variable.definition_period != periods.ETERNITY:
+        if self.variable.definition_period != DateUnit.ETERNITY:
             if period is None:
-                raise ValueError('A period must be specified to set values, except for variables with periods.ETERNITY as as period_definition.')
+                raise ValueError('A period must be specified to set values, except for variables with DateUnit.ETERNITY as as period_definition.')
             if (self.variable.definition_period != period.unit or period.size > 1):
                 name = self.variable.name
                 period_size_adj = f'{period.unit}' if (period.size == 1) else f'{period.size}-{period.unit}s'

openfisca_core/periods/__init__.py

@@ -42,5 +42,6 @@
     unit_weight,
     )
 
+from .date_unit import DateUnit  # noqa: F401
 from .instant_ import Instant  # noqa: F401
 from .period_ import Period  # noqa: F401

openfisca_core/periods/config.py

@@ -1,15 +1,17 @@
 import re
-import typing
+from typing import Dict
 
-DAY = 'day'
-MONTH = 'month'
-YEAR = 'year'
-ETERNITY = 'eternity'
+from .date_unit import DateUnit
+
+DAY = DateUnit.DAY
+MONTH = DateUnit.MONTH
+YEAR = DateUnit.YEAR
+ETERNITY = DateUnit.ETERNITY
 
 # Matches "2015", "2015-01", "2015-01-01"
 # Does not match "2015-13", "2015-12-32"
 INSTANT_PATTERN = re.compile(r"^\d{4}(-(0[1-9]|1[012]))?(-(0[1-9]|1[012])-(0[1-9]|[12][0-9]|3[01]))?$")
 
-date_by_instant_cache: typing.Dict = {}
-str_by_instant_cache: typing.Dict = {}
+date_by_instant_cache: Dict = {}
+str_by_instant_cache: Dict = {}
 year_or_month_or_day_re = re.compile(r'(18|19|20)\d{2}(-(0?[1-9]|1[0-2])(-([0-2]?\d|3[0-1]))?)?$')

openfisca_core/periods/date_unit.py

@@ -0,0 +1,54 @@
+from strenum import StrEnum
+
+
+class DateUnit(StrEnum):
+    """The date units of a rule system.
+
+    Examples:
+        >>> repr(DateUnit)
+        "<enum 'DateUnit'>"
+
+        >>> repr(DateUnit.DAY)
+        "<DateUnit.DAY: 'day'>"
+
+        >>> str(DateUnit.DAY)
+        'day'
+
+        >>> dict([(DateUnit.DAY, DateUnit.DAY.value)])
+        {<DateUnit.DAY: 'day'>: 'day'}
+
+        >>> list(DateUnit)
+        [<DateUnit.DAY: 'day'>, <DateUnit.MONTH: 'month'>, ...]
+
+        >>> len(DateUnit)
+        4
+
+        >>> DateUnit["DAY"]
+        <DateUnit.DAY: 'day'>
+
+        >>> DateUnit(DateUnit.DAY)
+        <DateUnit.DAY: 'day'>
+
+        >>> DateUnit.DAY in DateUnit
+        True
+
+        >>> "day" in list(DateUnit)
+        True
+
+        >>> DateUnit.DAY == "day"
+        True
+
+        >>> DateUnit.DAY.name
+        'DAY'
+
+        >>> DateUnit.DAY.value
+        'day'
+
+    .. versionadded:: 35.9.0
+
+    """
+
+    DAY = "day"
+    MONTH = "month"
+    YEAR = "year"
+    ETERNITY = "eternity"