Stop deep-copying the parameter tree in Reform.modify_parameters, and read in-place edits fresh

changelog.d/reform-modify-parameters-copy.fixed.md

@@ -0,0 +1 @@
+`Reform.modify_parameters` no longer deep-copies the baseline's parameter tree, so changes `apply()` made first are kept, and `TaxBenefitSystem.modify_parameters` drops every cached node at instant after the modifier runs, so in-place edits read fresh.

policyengine_core/parameters/parameter_node.py

@@ -235,6 +235,19 @@ def clear_parent_cache(self):
         if self.parent is not None:
             self.parent.clear_parent_cache()
 
+    def clear_descendant_caches(self) -> None:
+        """Drop the nodes at instant cached by this node and every node below it.
+
+        ``Parameter.update`` clears the nodes above it, as far as ``parent``
+        links reach. This reaches the nodes an edit made some other way leaves
+        behind: a change to ``values_list``, a scale's ``brackets`` or
+        ``children``.
+        """
+        self._at_instant_cache.clear()
+        for descendant in self.get_descendants():
+            if isinstance(descendant, ParameterNode):
+                descendant._at_instant_cache.clear()
+
     def mark_as_modified(self):
         self.modified = True
         if self.parent is not None:

policyengine_core/reforms/reform.py

@@ -1,7 +1,6 @@
 from __future__ import annotations
 
-import copy
-from typing import Callable, Union, TYPE_CHECKING
+from typing import Union, TYPE_CHECKING
 
 from policyengine_core.parameters import ParameterNode, Parameter
 from policyengine_core.taxbenefitsystems import TaxBenefitSystem
@@ -124,10 +123,8 @@ def __init__(self, baseline: TaxBenefitSystem):
         super().__init__(baseline.entities)
         self.baseline = baseline
         # Clone the parameter tree so in-place mutations inside ``apply()``
-        # (``set_parameter(...)`` or ``self.parameters.path.update(...)``)
-        # don't leak into the baseline. ``modify_parameters()`` was already
-        # deep-copying, but plenty of reforms mutate ``self.parameters``
-        # directly (bug C4).
+        # (``modify_parameters(...)`` or ``self.parameters.path.update(...)``)
+        # don't leak into the baseline (bug C4).
         if baseline.parameters is not None:
             self.parameters = baseline.parameters.clone()
         else:
@@ -158,29 +155,6 @@ def full_key(self) -> str:
             key = ".".join([baseline_full_key, key])
         return key
 
-    def modify_parameters(
-        self, modifier_function: Callable[[ParameterNode], ParameterNode]
-    ) -> None:
-        """Make modifications on the parameters of the legislation.
-
-        Call this function in `apply()` if the reform asks for legislation parameter modifications.
-
-        Args:
-            modifier_function: A function that takes a :obj:`.ParameterNode` and should return an object of the same type.
-        """
-        baseline_parameters = self.baseline.parameters
-        baseline_parameters_copy = copy.deepcopy(baseline_parameters)
-        reform_parameters = modifier_function(baseline_parameters_copy)
-        if not isinstance(reform_parameters, ParameterNode):
-            return ValueError(
-                "modifier_function {} in module {} must return a ParameterNode".format(
-                    modifier_function.__name__,
-                    modifier_function.__module__,
-                )
-            )
-        self.parameters = reform_parameters
-        self._parameters_at_instant_cache = {}
-
     @staticmethod
     def from_dict(
         parameter_values: dict,

policyengine_core/taxbenefitsystems/tax_benefit_system.py

@@ -796,33 +796,55 @@ def test(self, paths: str, verbose: bool = False) -> None:
 
     def modify_parameters(
         self, modifier_function: Callable[[ParameterNode], ParameterNode]
-    ) -> None:
+    ) -> TaxBenefitSystem | ValueError:
         """Make modifications on the parameters of the legislation.
 
         Call this function in `apply()` if the reform asks for legislation parameter modifications.
 
+        The modifier gets this system's own parameter tree, with every change
+        made to it so far, and may change it in place. A reform built with
+        ``Reform(baseline)`` holds a clone of the baseline's tree, so its changes
+        stay out of the baseline. A clone shares nested metadata and list values
+        with what it was cloned from, so a modifier replaces those rather than
+        editing them.
+
+        This system and its tree read as the changes left them, however the
+        modifier made them: every node cached at an instant is dropped once it
+        finishes, even if it raises. Edits it made before raising stay.
+
         Args:
             modifier_function: A function that takes a :obj:`.ParameterNode` and should return an object of the same type.
         """
-        if isinstance(modifier_function, dict):
-            for parameter in modifier_function:
-                for time_period in modifier_function[parameter]:
-                    param: Parameter = self.parameters.get_child(parameter)
-                    param.update(
-                        value=modifier_function[parameter][time_period],
-                        period=periods.period(time_period),
-                    )
-        else:
-            reform_parameters = modifier_function(self.parameters)
-            if not isinstance(reform_parameters, ParameterNode):
-                return ValueError(
-                    "modifier_function {} in module {} must return a ParameterNode".format(
-                        modifier_function.__name__,
-                        modifier_function.__module__,
+        parameters = self.parameters
+        reform_parameters = parameters
+        try:
+            if isinstance(modifier_function, dict):
+                for parameter in modifier_function:
+                    for time_period in modifier_function[parameter]:
+                        param: Parameter = self.parameters.get_child(parameter)
+                        param.update(
+                            value=modifier_function[parameter][time_period],
+                            period=periods.period(time_period),
+                        )
+            else:
+                reform_parameters = modifier_function(parameters)
+                if not isinstance(reform_parameters, ParameterNode):
+                    return ValueError(
+                        "modifier_function {} in module {} must return a ParameterNode".format(
+                            modifier_function.__name__,
+                            modifier_function.__module__,
+                        )
                     )
-                )
-            self.parameters = reform_parameters
-        self._parameters_at_instant_cache = {}
+                self.parameters = reform_parameters
+        finally:
+            self._parameters_at_instant_cache = {}
+            if parameters is not None:
+                parameters.clear_descendant_caches()
+            if (
+                isinstance(reform_parameters, ParameterNode)
+                and reform_parameters is not parameters
+            ):
+                reform_parameters.clear_descendant_caches()
         return self
 
     def add_modelled_policy_metadata(self):

tests/core/parameters/test_vectorial_parameter_node_copy.py

@@ -9,7 +9,7 @@
 ``recarray``.
 
 ``ParameterNodeAtInstant`` caches its vectorial node, and
-``Reform.modify_parameters`` deep-copies the baseline's parameter tree, so a
+``Reform.modify_parameters`` deep-copied the baseline's parameter tree, so a
 reform found a ``recarray`` where its cached node should be.
 """
 
@@ -200,7 +200,7 @@ def test_copied_parameter_tree_keeps_its_cached_vectorial_node(copier):
     assert_near(rate[families][housings].z2, expected)
 
 
-def test_reform_that_modifies_parameters_keeps_the_cached_vectorial_node(
+def test_reform_that_modifies_parameters_reads_the_vectorial_node(
     isolated_tax_benefit_system,
 ):
     system = isolated_tax_benefit_system
@@ -210,10 +210,11 @@ def test_reform_that_modifies_parameters_keeps_the_cached_vectorial_node(
 
     class KeepParameters(Reform):
         def apply(self):
-            # Deep-copies the baseline's tree, cache included.
+            # Gets the reform's own copy of the tree, not the baseline's, which
+            # has cached that node.
             self.modify_parameters(lambda parameters: parameters)
 
     general = KeepParameters(system).get_parameters_at_instant("2017-01-01").general
 
-    assert type(general._vectorial_node) is VectorialParameterNodeAtInstant
     assert_near(general[names], expected)
+    assert type(general._vectorial_node) is VectorialParameterNodeAtInstant