Roll back failed override batches and keep structural-reform nodes across rebuilds

changelog.d/2374.fixed.md

@@ -0,0 +1 @@
+- A failed batch of legacy parameter overrides is now rolled back whole, nodes a structural reform adds to the parameters survive later rebuilds and reprocessing, and the scenario docs say that `Scenario.apply` on a built simulation edits the annualised values.

docs/book/usage/scenarios.md

@@ -396,7 +396,9 @@ Each way of passing a reform reads period keys as it always has:
 
 A single-date key in a `reform=` dict overrides the annualised value on that day only. Formulas read annual values on 1 January, so `"2027-01-01"` sets fiscal year 2027 and `"2027-04-06"` changes nothing. The model logs a warning when a key can change nothing. Use `"start.stop"` keys for dated measures.
 
-Changes applied after the data is loaded (the default) leave the data as built, so a changed economic assumption does not re-uprate incomes. They still come before the steps that read parameters while the simulation is built: the Universal Credit rebalancing awards and the structural reforms that parameters switch on. Set `applied_before_data_load=True` on a scenario to project the data with changed economic assumptions too.
+Changes applied after the data is loaded (the default) leave the data as built, so a changed economic assumption does not re-uprate incomes. Dated changes (`parameter_changes` and `"start.stop"` or bare-year keys) still come before the steps that read parameters while the simulation is built: the Universal Credit rebalancing awards and the structural reforms that parameters switch on. Single-day keys and a scenario's own simulation modifier run after those steps, as they always have. Set `applied_before_data_load=True` on a scenario to project the data with changed economic assumptions too.
+
+All of this applies when the scenario or reform is passed to `Simulation(...)`. `Scenario.apply(simulation)` on a simulation that is already built edits its annualised values instead, as it always has: a `"2027-04-06.2028-04-05"` key there overrides those dates of the fiscal-year values, so it does not set fiscal year 2027, and the model logs a warning saying so. Pass the scenario when constructing the simulation for dated measures.
 
 ### Creating complex scenarios with simulation modifiers
 

policyengine_uk/tax_benefit_system.py

@@ -53,6 +53,7 @@
     Timeline,
     apply_legacy_overrides,
     copy_timelines,
+    graft_missing_nodes,
     requires_projection,
     snapshot_timelines,
     tree_defines,
@@ -173,8 +174,21 @@ def reset_parameter_caches(self):
         self.parameters._at_instant_cache = {}
 
     def reset_parameters(self) -> None:
-        """Reset parameters by reloading from the parameters directory."""
+        """Reset parameters by reloading from the parameters directory.
+
+        The processed parameters it replaces are kept until
+        ``process_parameters`` runs, so nodes a structural reform added to
+        them survive the reload. Only a processed tree is kept, and a second
+        reset before processing keeps the first one.
+        """
         self._parameters_at_instant_cache = {}
+        current = getattr(self, "parameters", None)
+        if (
+            getattr(self, "_parameters_before_reset", None) is None
+            and current is not None
+            and "baseline" in current.children
+        ):
+            self._parameters_before_reset = current
         self.load_parameters(self.parameters_dir)
         self.parameters = add_removed_parameter_aliases(self.parameters)
 
@@ -188,14 +202,34 @@ def process_parameters(self) -> None:
         The system's dated reform changes (``dated_assignments``) are written
         onto the parameters first, so a modifier that reloads and edits the
         parameters before processing them again (as ``no_economic_assumptions``
-        does) keeps the reform. Its legacy overrides of the annualised
-        parameters are then applied again. The parameters as they stood
-        before the changes are kept as the base for later rebuilds.
+        does) keeps the reform. Nodes a structural reform added to the
+        processed parameters before ``reset_parameters`` are added back (so
+        a node a modifier deletes after the reset comes back too), and the
+        legacy overrides of the annualised parameters are applied again.
+        The parameters as they stood before the changes are kept as the base
+        for later rebuilds. If processing fails, the system is left as it
+        was.
         """
+        previous = (self.parameters, self._projected_timelines, self._raw_parameters)
+        before_reset = getattr(self, "_parameters_before_reset", None)
         self._parameters_at_instant_cache = {}
-        self._rebuild(self.parameters, statutory=False)
-        self._apply_legacy_overrides()
-        self.reset_parameter_caches()
+        try:
+            # Processing changes the tree it is given, so it gets a copy: if
+            # it fails, the system keeps the tree it had.
+            self._rebuild(self.parameters.clone(), statutory=False)
+            if before_reset is not None:
+                graft_missing_nodes(before_reset, self.parameters)
+            self._apply_legacy_overrides()
+        except Exception:
+            (
+                self.parameters,
+                self._projected_timelines,
+                self._raw_parameters,
+            ) = previous
+            raise
+        finally:
+            self.reset_parameter_caches()
+        self._parameters_before_reset = None
 
     def apply_parameter_assignments(
         self, assignments: List[ParameterAssignment]
@@ -213,10 +247,10 @@ def apply_parameter_assignments(
         parameter takes under the reform's other changes (see
         ``policyengine_uk.utils.dated_reforms``).
 
-        A rebuild keeps the system's dated changes and legacy overrides. It
-        does not keep other edits made to the processed parameters, such as
-        a structural reform's, as ``reset_parameters`` never has. If the
-        rebuild fails, the system is left as it was.
+        A rebuild keeps the system's dated changes and legacy overrides, and
+        the nodes a structural reform added to the processed parameters. It
+        does not keep other edits made directly to existing processed
+        parameters. If the rebuild fails, the system is left as it was.
         """
         if not assignments:
             return
@@ -240,6 +274,8 @@ def apply_parameter_assignments(
             cached = _rebuilt_parameters_cache.get(key) if key is not None else None
             if cached is not None:
                 self.parameters = cached[0].clone()
+                # Shared with the memo: read only through the copying
+                # ``projected_timelines`` property.
                 self._projected_timelines = cached[1]
             else:
                 raw = (
@@ -254,6 +290,7 @@ def apply_parameter_assignments(
                         self.parameters.clone(),
                         copy_timelines(self._projected_timelines),
                     )
+            graft_missing_nodes(previous[1], self.parameters)
             self._apply_legacy_overrides()
         except Exception:
             (
@@ -270,12 +307,31 @@ def add_legacy_overrides(self, overrides: List[LegacyOverride]) -> None:
         """Apply changes to the annualised parameters and keep them.
 
         Kept so a later rebuild of the parameters applies them again, after
-        the dated changes. Overrides that fail to apply are not kept.
+        the dated changes. If any override fails, none is applied or kept.
         """
         if not overrides:
             return
+        # Every target is found before anything changes; ``update`` replaces
+        # a parameter's ``values_list`` and marks it and its ancestors
+        # modified, so keeping each list and those flags is enough to undo
+        # the batch.
+        targets = [self.parameters.get_child(override.path) for override in overrides]
+        saved = [(target, target.values_list) for target in targets]
+        flags = {}
+        for target in targets:
+            node = target
+            while node is not None and id(node) not in flags:
+                flags[id(node)] = (node, getattr(node, "modified", False))
+                node = getattr(node, "parent", None)
         try:
             apply_legacy_overrides(self.parameters, overrides)
+        except Exception:
+            for target, values_list in saved:
+                target.values_list = values_list
+                target.parent.clear_parent_cache()
+            for node, modified in flags.values():
+                node.modified = modified
+            raise
         finally:
             self.reset_parameter_caches()
         self.legacy_overrides = self.legacy_overrides + list(overrides)
@@ -352,6 +408,9 @@ def __init__(self):
         # when it has been processed itself (None: current law's).
         self._raw_parameters: Optional[ParameterNode] = None
         self._projected_timelines: Optional[Dict[str, Timeline]] = None
+        # The processed parameters ``reset_parameters`` replaced, until
+        # ``process_parameters`` adds back the nodes only they had.
+        self._parameters_before_reset: Optional[ParameterNode] = None
 
         # Set up and process parameters
         self.parameters_dir = COUNTRY_DIR / "parameters"
@@ -400,7 +459,11 @@ def statutory_parameter_defines() -> Callable[[str], bool]:
 
 
 def projected_timelines() -> Dict[str, Timeline]:
-    """Current law's projected timeline for each gov parameter, by name."""
+    """Current law's projected timeline for each gov parameter, by name.
+
+    Shared: do not change them. ``CountryTaxBenefitSystem.projected_timelines``
+    returns a copy.
+    """
     _ensure_caches()
     return _projected_timelines_cache
 

policyengine_uk/tests/test_dated_reform_routes.py

@@ -1024,3 +1024,210 @@ def test_preserve_calendar_dates_changes_keep_their_dates(start, length, value):
     while day < datetime.date(2035, 1, 1):
         assert parameter(day.isoformat()) == timeline_value(expected, day), day
         day += datetime.timedelta(days=37)
+
+
+# Review round 3 (GPT-6.1 Sol): failed override batches roll back, and
+# nodes a structural reform adds survive rebuilds and reprocessing.
+
+
+class _add_bonus(Reform):
+    def apply(self):
+        from policyengine_core.parameters import ParameterNode
+
+        self.parameters.add_child(
+            "test_budget",
+            ParameterNode("test_budget", data={"bonus": {"values": {"2000-01-01": 1}}}),
+        )
+
+
+def test_a_failed_override_batch_applies_none_of_its_overrides():
+    simulation = Simulation(situation=SITUATION)
+    system = simulation.tax_benefit_system
+    parameters = system.parameters
+    with pytest.raises(ValueError):
+        Scenario.from_reform(
+            {BASIC_RATE: {"2027": 0.3}, "gov.not_a_real_parameter": {"2027": 1}}
+        ).apply(simulation)
+    assert annual(simulation, BASIC_RATE, 2027) == 0.20
+    assert system.legacy_overrides == []
+    # An unknown path is refused before anything changes, and the tree is
+    # the one the system had.
+    assert system.parameters is parameters
+    Scenario.from_reform({BASIC_RATE: {"2027": 0.3}}).apply(simulation)
+    assert system.parameters is parameters
+    assert annual(simulation, BASIC_RATE, 2027) == 0.3
+    simulation.apply_parameter_changes({UC_SINGLE_OLD: {"2027": 500}})
+    assert annual(simulation, BASIC_RATE, 2027) == 0.3
+    assert annual(simulation, UC_SINGLE_OLD, 2027) == 500
+
+
+def test_a_batch_failing_after_an_update_restores_values_and_flags(monkeypatch):
+    simulation = Simulation(situation=SITUATION)
+    system = simulation.tax_benefit_system
+    rate = system.parameters.get_child(BASIC_RATE)
+    values_list = rate.values_list
+    # Processing marks parameters modified, so clear the flags on the path
+    # to the rate to see whether a failed batch leaves them set.
+    path, node = [], rate
+    while node is not None:
+        node.modified = False
+        path.append(node)
+        node = getattr(node, "parent", None)
+    # Reform.from_dict keys reach the overrides unparsed, so the second
+    # override fails after the first has updated the rate.
+    reform = Reform.from_dict(
+        {BASIC_RATE: {"2027-01-01": 0.3}, UC_SINGLE_OLD: {"not-a-date": 1}},
+        country_id="uk",
+    )
+    # Guard against the test going empty: the batch must fail inside the
+    # overrides, after the rate was updated, so the rollback runs.
+    original = tbs_module.apply_legacy_overrides
+    seen = {}
+
+    def spy(parameters, overrides):
+        try:
+            return original(parameters, overrides)
+        except Exception:
+            seen["rate_updated"] = rate.values_list is not values_list
+            raise
+
+    monkeypatch.setattr(tbs_module, "apply_legacy_overrides", spy)
+    with pytest.raises(ValueError):
+        Scenario.from_reform(reform).apply(simulation)
+    assert seen == {"rate_updated": True}
+    assert rate.values_list is values_list
+    assert annual(simulation, BASIC_RATE, 2027) == 0.20
+    assert system.legacy_overrides == []
+    assert not any(node.modified for node in path)
+
+
+def test_a_rebuild_keeps_a_structural_reforms_parameter():
+    simulation = Simulation(
+        situation=SITUATION, reform=(_add_bonus, {"test_budget.bonus": 7})
+    )
+    simulation.apply_parameter_changes({BASIC_RATE: {"2027": 0.21}})
+    parameters = simulation.tax_benefit_system.parameters
+    assert parameters.test_budget.bonus("2027-01-01") == 7
+    assert annual(simulation, BASIC_RATE, 2027) == 0.21
+    # Repeated rebuilds keep it too.
+    simulation.apply_parameter_changes({UC_SINGLE_OLD: {"2027": 500}})
+    assert simulation.tax_benefit_system.parameters.test_budget.bonus("2027-01-01") == 7
+
+
+def test_reprocessing_keeps_a_structural_reforms_parameter():
+    from policyengine_uk.scenarios.no_economic_assumptions import (
+        no_economic_assumptions,
+    )
+
+    simulation = Simulation(
+        situation=SITUATION, reform=(_add_bonus, {"test_budget.bonus": 7})
+    )
+    no_economic_assumptions.apply(simulation)
+    parameters = simulation.tax_benefit_system.parameters
+    assert parameters.test_budget.bonus("2027-01-01") == 7
+    assert annual(simulation, PERSONAL_ALLOWANCE, 2035) == 12_570
+    assert simulation.tax_benefit_system._parameters_before_reset is None
+
+
+def test_failed_reprocessing_leaves_the_system_as_it_was(monkeypatch):
+    simulation = Simulation(situation=SITUATION)
+    Scenario.from_reform({BASIC_RATE: {"2027": 0.3}}).apply(simulation)
+    system = simulation.tax_benefit_system
+    system.reset_parameters()
+    reloaded = (system.parameters, system._projected_timelines, system._raw_parameters)
+
+    # Fail after the rebuild has replaced the parameters, while replaying
+    # the legacy override.
+    def fail(*args, **kwargs):
+        raise RuntimeError("replay failed")
+
+    monkeypatch.setattr(tbs_module, "apply_legacy_overrides", fail)
+    with pytest.raises(RuntimeError, match="replay failed"):
+        system.process_parameters()
+    assert system.parameters is reloaded[0]
+    assert system._projected_timelines is reloaded[1]
+    assert system._raw_parameters is reloaded[2]
+    # The tree is as reloaded, not processed: no baseline child, and the
+    # basic rate keeps its legal dates rather than fiscal-year values.
+    assert "baseline" not in system.parameters.children
+    # Processing again then works.
+    monkeypatch.undo()
+    system.process_parameters()
+    assert "baseline" in system.parameters.children
+    assert annual(simulation, BASIC_RATE, 2027) == 0.3
+
+
+class _add_top_band(Reform):
+    def apply(self):
+        scale = self.parameters.gov.hmrc.income_tax.rates.uk
+        band = scale.brackets[-1].clone()
+        band.parent = scale
+        scale.brackets.append(band)
+        band.threshold.update(period="year:2000:200", value=200_000)
+        band.rate.update(period="year:2000:200", value=0.5)
+
+
+def test_a_rebuild_keeps_a_band_a_structural_reform_appended():
+    simulation = Simulation(situation=SITUATION, reform=_add_top_band)
+    # The scale has four bands, uk[0] to uk[3]; the reform adds uk[4].
+    top = "gov.hmrc.income_tax.rates.uk[4].rate"
+    assert annual(simulation, top, 2027) == 0.5
+    simulation.apply_parameter_changes({BASIC_RATE: {"2027": 0.21}})
+    assert annual(simulation, top, 2027) == 0.5
+    assert annual(simulation, BASIC_RATE, 2027) == 0.21
+
+
+def test_a_second_reset_keeps_the_processed_tree_for_grafting():
+    simulation = Simulation(
+        situation=SITUATION, reform=(_add_bonus, {"test_budget.bonus": 7})
+    )
+    system = simulation.tax_benefit_system
+    system.reset_parameters()
+    system.reset_parameters()
+    system.process_parameters()
+    assert system.parameters.test_budget.bonus("2027-01-01") == 7
+
+
+def test_grafting_adds_only_missing_nodes():
+    from policyengine_core.parameters import ParameterNode
+
+    from policyengine_uk.utils.dated_reforms import graft_missing_nodes
+
+    source = ParameterNode(
+        "",
+        data={
+            "a": {
+                "x": {"values": {"2000-01-01": 1}},
+                "y": {"values": {"2000-01-01": 2}},
+            },
+            "baseline": {"z": {"values": {"2000-01-01": 3}}},
+        },
+    )
+    target = ParameterNode("", data={"a": {"x": {"values": {"2000-01-01": 10}}}})
+    added = graft_missing_nodes(source, target)
+    assert added == ["a.y"]
+    assert target.a.x("2001-01-01") == 10
+    assert target.a.y("2001-01-01") == 2
+    assert "baseline" not in target.children
+
+
+def test_grafting_appends_missing_scale_brackets():
+    from policyengine_core.parameters import ParameterNode
+
+    from policyengine_uk.utils.dated_reforms import graft_missing_nodes
+
+    def band(threshold, rate):
+        return {
+            "threshold": {"values": {"2000-01-01": threshold}},
+            "rate": {"values": {"2000-01-01": rate}},
+        }
+
+    source = ParameterNode(
+        "", data={"s": {"brackets": [band(0, 0.1), band(10, 0.2), band(20, 0.3)]}}
+    )
+    target = ParameterNode("", data={"s": {"brackets": [band(0, 0.15), band(10, 0.2)]}})
+    assert graft_missing_nodes(source, target) == ["s[2]"]
+    assert len(target.s.brackets) == 3
+    assert target.s.brackets[0].rate("2001-01-01") == 0.15
+    assert target.s.brackets[2].rate("2001-01-01") == 0.3
+    assert target.s.brackets[2].parent is target.s