Restore inputs as inputs, evict the fast cache on holder writes, reject non-numeric uprating and defined_for

changelog.d/fix-restore-registry-holder-cache-enum.changed.md

@@ -0,0 +1 @@
+Restore non-derived dumped values as recorded inputs using derived-period marks and an explicit format marker, warn for ambiguous legacy dumps, evict affected fast-cache entries on holder writes and deletes, reject non-numeric uprating and defined_for conditions, and preserve existing variables when a replacement is rejected.

policyengine_core/holders/holder.py

@@ -109,6 +109,7 @@ def delete_arrays(
         self._memory_storage.delete(period, branch_name)
         if self._disk_storage:
             self._disk_storage.delete(period, branch_name)
+        self._evict_fast_cache(period, branch_name, contained=True)
 
     def _get_array_from_storage(
         self, period: Period, branch_name: str = "default"
@@ -380,6 +381,7 @@ def _set(
             self._disk_storage.put(value, period, branch_name, derived=derived)
         else:
             self._memory_storage.put(value, period, branch_name, derived=derived)
+        self._evict_fast_cache(period, branch_name)
         if user_input_contexts:
             if not hasattr(simulation, "_user_input_keys"):
                 simulation._user_input_keys = set()
@@ -507,3 +509,52 @@ def get_input_periods(self, branch_name: str = "default") -> List[Period]:
             for period, (_, storage, stored_branch) in read.items()
             if not storage.is_derived(period, stored_branch)
         ]
+
+    def _evict_fast_cache(
+        self, period: Period, branch_name: str, contained: bool = False
+    ) -> None:
+        """Drop the simulation's ``_fast_cache`` entries a storage write or delete makes stale.
+
+        ``Simulation.calculate`` answers a repeated request from
+        ``_fast_cache``, keyed by ``(variable name, requested period)``,
+        before it reads this holder. So every write into, or delete from,
+        this holder's storage drops the entries for the periods it changes;
+        otherwise ``calculate`` kept returning the value the storage no
+        longer held (for example after ``holder.set_input``).
+
+        The fast cache belongs to this holder's simulation: a branch has its
+        own holders and its own fast cache, and keeps the values it started
+        with, so nothing outside this simulation is touched. Nor is anything
+        here when ``branch_name`` is a branch this simulation does not read.
+
+        A write changes one storage key: ``period``, or, for an ETERNITY
+        variable, the one value every period reads. A delete (``contained``)
+        removes every period ``period`` contains, or every period when
+        ``period`` is ``None``.
+        """
+        simulation = self.simulation
+        fast_cache = getattr(simulation, "_fast_cache", None)
+        if not fast_cache:
+            return
+        name = self.variable.name
+        drop_all = period is None or self.variable.definition_period == periods.ETERNITY
+        if not drop_all:
+            period = periods.period(period)
+            if not contained and (name, period) not in fast_cache:
+                return
+        visible_branch_names = getattr(simulation, "_get_visible_branch_names", None)
+        if visible_branch_names is not None and branch_name not in (
+            visible_branch_names()
+        ):
+            return
+        if not drop_all and not contained:
+            del fast_cache[(name, period)]
+            return
+        stale_keys = [
+            key
+            for key in fast_cache
+            if key[0] == name
+            and (drop_all or not isinstance(key[1], Period) or period.contains(key[1]))
+        ]
+        for key in stale_keys:
+            del fast_cache[key]

policyengine_core/simulations/simulation.py

@@ -1063,6 +1063,17 @@ def _calculate(self, variable_name: str, period: Period = None) -> ArrayLike:
         self._check_period_consistency(period, variable)
 
         if variable.defined_for is not None:
+            # Registration rejects a non-numeric defined_for variable (see
+            # ``TaxBenefitSystem._check_defined_for``). This catches one set,
+            # or a variable replaced, afterwards. It reads the variable's
+            # type, not the values: mapped to a group entity, an Enum's
+            # indices are summed into numbers, and str or date values fail
+            # inside the mapping.
+            defined_for_variable = self.tax_benefit_system.get_variable(
+                variable.defined_for
+            )
+            if defined_for_variable is not None:
+                variable.check_defined_for_variable(defined_for_variable)
             mask = (
                 self.calculate(variable.defined_for, period, map_to=variable.entity.key)
                 > 0
@@ -1109,6 +1120,9 @@ def _calculate(self, variable_name: str, period: Period = None) -> ArrayLike:
                         if known_period in input_periods
                     ]
                 if earlier_input_periods:
+                    # Registration rejects these; an ``uprating`` assigned
+                    # past the setter gets the same message here.
+                    variable.check_uprating_value_type()
                     # Take the latest period from the filtered list itself.
                     # Indexing ``known_periods`` with a position in the
                     # filtered list picked the wrong period whenever a later

policyengine_core/taxbenefitsystems/tax_benefit_system.py

@@ -51,6 +51,7 @@
 from policyengine_core.periods import Instant, Period
 from policyengine_core.populations import GroupPopulation, Population
 from policyengine_core.variables import Variable
+from policyengine_core.variables.variable import NUMERIC_VALUE_TYPES
 
 log = logging.getLogger(__name__)
 
@@ -100,6 +101,8 @@ class TaxBenefitSystem:
     """Short list of basic inputs to get medium accuracy."""
     modelled_policies: str = None
     """A YAML filepath containing metadata describing the modelled policies."""
+    _defined_for_checks_deferred: bool = False
+    """Whether ``load_variable`` leaves ``defined_for`` checks to a later pass over every variable."""
 
     def __init__(self, entities: Sequence[Entity] = None, reform=None) -> None:
         if entities is None:
@@ -127,7 +130,15 @@ def __init__(self, entities: Sequence[Entity] = None, reform=None) -> None:
         self.variable_module_metadata = {}
 
         if self.variables_dir is not None:
-            self.add_variables_from_directory(self.variables_dir)
+            # Check every variable once, after the whole directory is loaded,
+            # instead of scanning the variables loaded so far each time a
+            # non-numeric one is added.
+            self._defined_for_checks_deferred = True
+            try:
+                self.add_variables_from_directory(self.variables_dir)
+            finally:
+                self._defined_for_checks_deferred = False
+            self._check_defined_for_variables()
         self.data_modified = False
 
         if self.parameters_dir is not None:
@@ -223,10 +234,44 @@ def load_variable(
             )
 
         variable = variable_class(baseline_variable=baseline_variable)
+        if not self._defined_for_checks_deferred:
+            self._check_defined_for(variable)
         self.variables[variable.name] = variable
 
         return variable
 
+    def _check_defined_for(self, variable: Variable) -> None:
+        """Check the ``defined_for`` links that ``variable`` takes part in.
+
+        Both directions: the variable ``variable`` is defined for, if it is
+        registered, and, when ``variable`` cannot be compared with zero, every
+        registered variable defined for it. See
+        :meth:`Variable.check_defined_for_variable`.
+        """
+        defined_for = variable.defined_for
+        if defined_for is not None:
+            defined_for_variable = (
+                variable
+                if defined_for == variable.name
+                else self.variables.get(defined_for)
+            )
+            if defined_for_variable is not None:
+                variable.check_defined_for_variable(defined_for_variable)
+        if variable.value_type in NUMERIC_VALUE_TYPES:
+            return
+        for other in self.variables.values():
+            if other.defined_for == variable.name and other.name != variable.name:
+                other.check_defined_for_variable(variable)
+
+    def _check_defined_for_variables(self) -> None:
+        """Check every registered variable's ``defined_for`` variable."""
+        for variable in self.variables.values():
+            if variable.defined_for is None:
+                continue
+            defined_for_variable = self.variables.get(variable.defined_for)
+            if defined_for_variable is not None:
+                variable.check_defined_for_variable(defined_for_variable)
+
     def add_variable(self, variable: Type[Variable]) -> Variable:
         """Adds an OpenFisca variable to the tax and benefit system.
 
@@ -251,9 +296,14 @@ def replace_variable(self, variable: Type[Variable]) -> None:
         :param Variable variable: New variable class to add. Must be a subclass of Variable.
         """
         name = variable.__name__
-        if self.variables.get(name) is not None:
-            del self.variables[name]
-        self.load_variable(variable, update=False)
+        replaced = self.variables.pop(name, None)
+        try:
+            self.load_variable(variable, update=False)
+        except Exception:
+            # A rejected replacement leaves the system as it was.
+            if replaced is not None:
+                self.variables[name] = replaced
+            raise
         self.data_modified = True
 
     def update_variable(self, variable: Type[Variable]) -> Variable:

policyengine_core/tools/simulation_dumper.py

@@ -2,21 +2,37 @@
 
 
 import os
+import warnings
 
 import numpy as np
 
 from policyengine_core.data_storage import OnDiskStorage
 from policyengine_core.periods import ETERNITY
 from policyengine_core.simulations import Simulation
 
-# Periods, one per line, whose dumped value the simulation calculated (see
-# ``Holder.is_derived``), so a restored simulation does not carry them over.
+# Next to each variable's arrays: the periods, one per line, whose dumped
+# value the simulation calculated (see ``Holder.is_derived``). The file is
+# left out when there are none. ``restore_simulation`` restores these values
+# as calculated ones, which auto-carry-over and uprating never read and
+# ``apply_reform`` drops, and every other value as an input.
 DERIVED_PERIODS_FILE = "derived_periods.txt"
 
+# In ``__entities__``, written last by every dump: it says the dump lists
+# every calculated value in ``DERIVED_PERIODS_FILE``, so a variable without
+# that file has none. Dumps written before policyengine-core 3.32.16 list no
+# calculated value, whatever they hold, and those written before this file
+# existed have no marker. (``restore_simulation`` reads only named files in
+# ``__entities__``, so earlier versions restore a dump with the marker.)
+DUMP_FORMAT_FILE = "dump_format.txt"
+DUMP_FORMAT_VERSION = 1
+
 
 def dump_simulation(simulation, directory):
     """
     Write simulation data to directory, so that it can be restored later.
+
+    Writes the default branch's values, listing the periods of those the
+    simulation calculated (see ``restore_simulation``).
     """
     parent_directory = os.path.abspath(os.path.join(directory, os.pardir))
     if not os.path.isdir(parent_directory):  # To deal with reforms
@@ -38,10 +54,32 @@ def dump_simulation(simulation, directory):
         for holder in entity._holders.values():
             _dump_holder(holder, directory)
 
+    with open(os.path.join(entities_dump_dir, DUMP_FORMAT_FILE), "w") as file:
+        file.write(f"{DUMP_FORMAT_VERSION}\n")
+
 
 def restore_simulation(directory, tax_benefit_system, **kwargs):
     """
     Restore simulation from directory
+
+    Each value is restored as the dumped simulation stored it: one it
+    calculated (listed in the variable's ``derived_periods.txt``) as a
+    calculated value, and every other value as an input, recorded in
+    ``_user_input_keys`` as ``set_input`` records one, so ``apply_reform``
+    keeps the inputs and drops the calculated values, as it would have in
+    the dumped simulation.
+
+    A dump lists which values were calculated, not where an input came
+    from: a value the dumped simulation stored with ``put_in_cache``
+    without ``derived=True`` (an input to auto-carry-over and uprating that
+    ``set_input`` did not record, so that ``apply_reform`` dropped it there)
+    is restored as a recorded input too.
+
+    A dump written before policyengine-core 3.32.16 lists no calculated
+    value, so every value in it is restored as an input. A dump without a
+    format marker that lists none may be one of those, so restoring it
+    warns: ``apply_reform`` then keeps any calculated value in it as dumped
+    instead of recalculating it.
     """
     simulation = Simulation(
         tax_benefit_system, tax_benefit_system.instantiate_entities()
@@ -59,11 +97,26 @@ def restore_simulation(directory, tax_benefit_system, **kwargs):
         _restore_entity(population, entities_dump_dir)
         population.count = person_count
 
-    variables_to_restore = (
+    variables_to_restore = [
         variable for variable in os.listdir(directory) if variable != "__entities__"
+    ]
+    lists_calculated_values = os.path.exists(
+        os.path.join(entities_dump_dir, DUMP_FORMAT_FILE)
     )
     for variable in variables_to_restore:
-        _restore_holder(simulation, variable, directory)
+        if _restore_holder(simulation, variable, directory):
+            lists_calculated_values = True
+    if variables_to_restore and not lists_calculated_values:
+        warnings.warn(
+            f"The simulation dump in {directory} does not say which of its "
+            f"values were calculated: it has no {DUMP_FORMAT_FILE} and lists "
+            "no calculated value, as dumps written by policyengine-core "
+            "before 3.32.16 never do. Every value in it is restored as an "
+            "input, so apply_reform keeps any calculated value in it as "
+            "dumped instead of recalculating it. Dump the simulation again "
+            "to record which values were calculated.",
+            stacklevel=2,
+        )
 
     return simulation
 
@@ -136,6 +189,9 @@ def _restore_entity(population, directory):
 
 
 def _restore_holder(simulation, variable, directory):
+    """Restore one variable's values: those its ``DERIVED_PERIODS_FILE``
+    lists as calculated values, every other one as an input. Return whether
+    it lists any."""
     storage_dir = os.path.join(directory, variable)
     is_variable_eternal = (
         simulation.tax_benefit_system.get_variable(variable).definition_period
@@ -156,4 +212,23 @@ def _restore_holder(simulation, variable, directory):
 
     for period in disk_storage.get_known_periods():
         value = disk_storage.get(period)
-        holder.put_in_cache(value, period, derived=str(period) in derived_periods)
+        if str(period) in derived_periods:
+            holder.put_in_cache(value, period, derived=True)
+        else:
+            _restore_input(simulation, holder, period, value)
+    return bool(derived_periods)
+
+
+def _restore_input(simulation, holder, period, value):
+    """Store ``value`` as an input, recorded as ``set_input`` records one.
+
+    ``Holder.set_input`` would also run the variable's ``set_input`` helper,
+    but a dump holds values already split into the variable's own periods.
+    """
+    if not hasattr(simulation, "_user_input_contexts"):
+        simulation._user_input_contexts = []
+    simulation._user_input_contexts.append(simulation.branch_name)
+    try:
+        holder._set(period, value, simulation.branch_name)
+    finally:
+        simulation._user_input_contexts.pop()