From ea99d5d5f78d71ae1373b3064b58f8b1a2a43470 Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Sun, 1 Nov 2026 16:25:58 -0500 Subject: [PATCH 1/2] Stop deep-copying the parameter tree in Reform.modify_parameters Reform.__init__ already clones the baseline's tree, so the deep copy in Reform.modify_parameters paid for a second full copy, threw away what apply() had changed, and carried the baseline's cached nodes into the reform. Reform now inherits TaxBenefitSystem.modify_parameters. That method now drops every node cached at an instant once the modifier finishes, in the tree it handed over and any tree it got back, so edits made without Parameter.update read fresh on both paths a reform runs. Fixes #618 Co-Authored-By: Claude Opus 5.5 --- .../reform-modify-parameters-copy.fixed.md | 1 + .../parameters/parameter_node.py | 13 + policyengine_core/reforms/reform.py | 32 +- .../taxbenefitsystems/tax_benefit_system.py | 58 ++- .../test_vectorial_parameter_node_copy.py | 9 +- tests/core/test_reform_modify_parameters.py | 442 ++++++++++++++++++ .../test_reform_modify_parameters_property.py | 213 +++++++++ tests/core/test_reforms.py | 1 - .../core/test_tracing_parameter_isolation.py | 7 +- tests/fixtures/reform_modifiers.py | 247 ++++++++++ 10 files changed, 968 insertions(+), 55 deletions(-) create mode 100644 changelog.d/reform-modify-parameters-copy.fixed.md create mode 100644 tests/core/test_reform_modify_parameters.py create mode 100644 tests/core/test_reform_modify_parameters_property.py create mode 100644 tests/fixtures/reform_modifiers.py diff --git a/changelog.d/reform-modify-parameters-copy.fixed.md b/changelog.d/reform-modify-parameters-copy.fixed.md new file mode 100644 index 000000000..18e82ecbf --- /dev/null +++ b/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. diff --git a/policyengine_core/parameters/parameter_node.py b/policyengine_core/parameters/parameter_node.py index 8338716f2..c807a37ad 100644 --- a/policyengine_core/parameters/parameter_node.py +++ b/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: diff --git a/policyengine_core/reforms/reform.py b/policyengine_core/reforms/reform.py index afcee0358..1b7298944 100644 --- a/policyengine_core/reforms/reform.py +++ b/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, diff --git a/policyengine_core/taxbenefitsystems/tax_benefit_system.py b/policyengine_core/taxbenefitsystems/tax_benefit_system.py index ccadd2a39..3dbf45447 100644 --- a/policyengine_core/taxbenefitsystems/tax_benefit_system.py +++ b/policyengine_core/taxbenefitsystems/tax_benefit_system.py @@ -801,28 +801,50 @@ def modify_parameters( 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. + 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): diff --git a/tests/core/parameters/test_vectorial_parameter_node_copy.py b/tests/core/parameters/test_vectorial_parameter_node_copy.py index 98af4f9b2..439f0de28 100644 --- a/tests/core/parameters/test_vectorial_parameter_node_copy.py +++ b/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 diff --git a/tests/core/test_reform_modify_parameters.py b/tests/core/test_reform_modify_parameters.py new file mode 100644 index 000000000..e471ddb5a --- /dev/null +++ b/tests/core/test_reform_modify_parameters.py @@ -0,0 +1,442 @@ +"""``modify_parameters`` gives a modifier the system's own tree and reads it fresh. + +A reform runs as an instance (``MyReform(baseline)``) or in place +(``system.apply_reform_set(MyReform)``). ``Reform.modify_parameters`` used to +deep-copy the baseline's tree for each call, which had two effects on the +instance path. It threw away what ``apply()`` had changed so far, and it carried +the baseline's cached nodes into the copy. A modifier that edited values, +brackets or children without ``Parameter.update`` then read the baseline's +values at any instant the baseline had cached. In place, the modifier gets the +system's own tree, with the same stale nodes. + +Both paths now run ``TaxBenefitSystem.modify_parameters``. See #618. + +``test_reform_modify_parameters_property.py`` has the properties. +""" + +import copy + +import pytest + +from policyengine_core.country_template import CountryTaxBenefitSystem +from policyengine_core.country_template.reforms.modify_social_security_taxation import ( + modify_social_security_taxation, +) +from policyengine_core.parameters import ParameterNode, ParameterScaleBracket +from policyengine_core.periods import instant +from policyengine_core.reforms import Reform, set_parameter +from tests.fixtures.reform_modifiers import ( + ModifierFailed, + assert_keeps_its_cached_nodes, + cached_nodes, + new_system, + read, + system_reads, + warm, +) + +INSTANT = "2017-01-01" + + +def as_instance(reform, system): + return reform(system) + + +def in_place(reform, system): + system.apply_reform_set(reform) + return system + + +BOTH_PATHS = pytest.mark.parametrize( + "run", [as_instance, in_place], ids=["as an instance", "in place"] +) + + +def social_security_thresholds(system): + return system.get_parameters_at_instant( + INSTANT + ).taxes.social_security_contribution.thresholds + + +def double_basic_income(parameters): + parameters.benefits.basic_income.update( + period=f"year:{INSTANT}:1", + value=2 * parameters.benefits.basic_income(INSTANT), + ) + return parameters + + +# ----- Stale values ----- # + + +class modify_brackets_in_place(Reform): + """The template reform, in the form ``apply_reform_set`` can run: there + ``self`` is the system, which has no ``modify_brackets``.""" + + def apply(self): + self.modify_parameters(modify_social_security_taxation.modify_brackets) + + +@pytest.mark.parametrize( + "run, reform", + [ + (as_instance, modify_social_security_taxation), + (in_place, modify_brackets_in_place), + ], + ids=["as an instance", "in place"], +) +@pytest.mark.parametrize("read_first", [True, False], ids=["read", "never read"]) +def test_template_reform_reads_its_modifier(run, reform, read_first): + system = CountryTaxBenefitSystem() + if read_first: + assert social_security_thresholds(system) == [0.0, 6000.0, 12400.0] + + reformed = run(reform, system) + + assert social_security_thresholds(reformed) == [6000.0, 12400.0, 40000] + + +def assign_a_value(parameters): + parameters.benefit.amount.values_list[0].value = -1 + + +def remove_a_value(parameters): + parameters.benefit.amount.values_list.pop(0) + + +def assign_a_brackets_value(parameters): + parameters.tax.brackets[1].rate.values_list[0].value = -1 + + +def delete_a_bracket(parameters): + del parameters.tax.brackets[0] + + +def append_a_bracket(parameters): + parameters.tax.brackets.append( + ParameterScaleBracket( + "new_bracket", + data={ + "threshold": {"2015-01-01": {"value": 5_000}}, + "rate": {"2015-01-01": {"value": 0.5}}, + }, + ) + ) + + +def add_a_child(parameters): + parameters.benefit.add_child("new", ParameterNode(data={"leaf": {"2015-01-01": 1}})) + + +def remove_a_child(parameters): + del parameters.benefit.children["eligible"] + + +def update_a_parameter_in_a_bracket(parameters): + """A loaded bracket has no ``parent``, so ``update`` stops at the bracket.""" + parameters.tax.brackets[1].rate.update(period="year:2018:1", value=-1) + + +@BOTH_PATHS +@pytest.mark.parametrize( + "edit", + [ + assign_a_value, + remove_a_value, + assign_a_brackets_value, + delete_a_bracket, + append_a_bracket, + add_a_child, + remove_a_child, + update_a_parameter_in_a_bracket, + ], +) +def test_edits_of_a_warm_tree_read_as_a_cold_tree_of_the_same_edits(run, edit): + def modify(parameters): + edit(parameters) + return parameters + + class edited(Reform): + def apply(self): + self.modify_parameters(modify) + + warmed = new_system() + warm(warmed) + reformed = run(edited, warmed) + + assert system_reads(reformed) == system_reads(run(edited, new_system())) + + +def test_dict_form_clears_what_an_update_under_a_bracket_leaves_behind(): + """A loaded bracket has no ``parent``, so ``update`` stops at the bracket.""" + system = new_system() + warm(system) + + system.modify_parameters({"tax[1].rate": {"year:2017:1": 0.9}}) + + assert system.get_parameters_at_instant(INSTANT).tax.rates == [0.1, 0.9] + assert system.parameters(INSTANT).tax.rates == [0.1, 0.9] + + +@BOTH_PATHS +@pytest.mark.parametrize( + "replace", [copy.deepcopy, ParameterNode.clone], ids=["deep copy", "clone"] +) +def test_a_tree_the_modifier_returns_reads_fresh(run, replace): + def modify(parameters): + parameters = replace(parameters) + parameters.benefit.amount.values_list[0].value = -1 + return parameters + + class replaces_the_tree(Reform): + def apply(self): + self.modify_parameters(modify) + + system = new_system() + warm(system) + + reformed = run(replaces_the_tree, system) + + assert reformed.parameters(INSTANT).benefit.amount == 100 + assert reformed.parameters("2020-01-01").benefit.amount == -1 + assert reformed.get_parameters_at_instant("2020-01-01").benefit.amount == -1 + + +def test_a_tree_the_modifier_edits_and_drops_reads_fresh(): + """Anything else holding that tree sees the edits as well.""" + system = new_system() + warm(system) + tree = system.parameters + + def modify(parameters): + parameters.benefit.amount.values_list[0].value = -1 + return ParameterNode(data={"leaf": {"2015-01-01": 1}}) + + system.modify_parameters(modify) + + assert tree("2020-01-01").benefit.amount == -1 + assert system.parameters("2020-01-01").leaf == 1 + + +def test_a_modifier_that_raises_leaves_a_tree_that_reads_fresh(): + """The exception reaches the caller, who may keep using the system.""" + system = new_system() + warm(system) + + def modify(parameters): + parameters.benefit.amount.values_list[0].value = -1 + raise ModifierFailed + + with pytest.raises(ModifierFailed): + system.modify_parameters(modify) + + assert system.get_parameters_at_instant("2020-01-01").benefit.amount == -1 + assert system.parameters("2020-01-01").benefit.amount == -1 + + +def test_clear_descendant_caches_reaches_every_node_below_and_none_above(): + system = new_system() + warm(system) + parameters = system.parameters + + parameters.benefit.clear_descendant_caches() + assert not parameters.benefit._at_instant_cache + assert parameters._at_instant_cache + + parameters.clear_descendant_caches() + assert not any( + node._at_instant_cache for node, _ in cached_nodes(parameters).values() + ) + + +# ----- Redundant copy, and the changes it dropped ----- # + + +@BOTH_PATHS +def test_changes_made_earlier_in_apply_are_kept(run): + class update_then_two_modifiers(Reform): + def apply(self): + self.parameters.taxes.income_tax_rate.update( + period=f"year:{INSTANT}:1", value=0.2 + ) + self.modify_parameters(double_basic_income) + self.modify_parameters(double_basic_income) + + reformed = run(update_then_two_modifiers, CountryTaxBenefitSystem()) + + at_instant = reformed.get_parameters_at_instant(INSTANT) + assert at_instant.taxes.income_tax_rate == 0.2 + assert at_instant.benefits.basic_income == 2400 + + +@BOTH_PATHS +def test_set_parameter_reforms_build_on_each_other(run): + first = set_parameter( + "taxes.income_tax_rate", 0.2, period="year:2017:1", return_modifier=True + ) + second = set_parameter( + "benefits.basic_income", 900, period="year:2017:1", return_modifier=True + ) + + class set_both(Reform): + def apply(self): + self.modify_parameters(first) + self.modify_parameters(second) + + at_instant = run(set_both, CountryTaxBenefitSystem()).get_parameters_at_instant( + INSTANT + ) + + assert at_instant.taxes.income_tax_rate == 0.2 + assert at_instant.benefits.basic_income == 900 + + +def test_a_reform_of_a_reform_keeps_the_first_reforms_changes(): + class first(Reform): + def apply(self): + self.modify_parameters(double_basic_income) + + class second(Reform): + def apply(self): + self.modify_parameters(double_basic_income) + + system = CountryTaxBenefitSystem() + once = first(system) + twice = second(once) + + assert system.parameters.benefits.basic_income(INSTANT) == 600 + assert once.parameters.benefits.basic_income(INSTANT) == 1200 + assert twice.parameters.benefits.basic_income(INSTANT) == 2400 + + +def test_a_reform_hands_its_modifier_the_tree_it_already_holds(monkeypatch): + system = new_system() + handed = [] + + def fail_to_deep_copy(self, memo): + raise AssertionError("A reform has no reason to deep-copy a parameter tree.") + + monkeypatch.setattr(ParameterNode, "__deepcopy__", fail_to_deep_copy, raising=False) + + class records_its_tree(Reform): + def apply(self): + tree = self.parameters + + def modify(parameters): + handed.append(parameters is tree) + return parameters + + self.modify_parameters(modify) + + assert records_its_tree(system).parameters is not system.parameters + assert handed == [True] + + +@BOTH_PATHS +def test_a_modifier_that_returns_no_tree_keeps_its_edits(run): + """It still gets ``ValueError`` back rather than raised: see #618.""" + results = [] + + def modify(parameters): + parameters.benefit.amount.update(period="year:2020:1", value=-1) + + class returns_nothing(Reform): + def apply(self): + results.append(self.modify_parameters(modify)) + + system = new_system() + warm(system) + + reformed = run(returns_nothing, system) + + assert type(results[0]) is ValueError + assert reformed.get_parameters_at_instant("2020-01-01").benefit.amount == -1 + + +def test_a_reform_instance_takes_the_dict_form(): + class dict_form(Reform): + def apply(self): + self.modify_parameters({"benefit.amount": {"year:2020:1": -1}}) + + system = new_system() + warm(system) + + assert dict_form(system).parameters("2020-01-01").benefit.amount == -1 + assert system.parameters("2020-01-01").benefit.amount == 140 + + +# ----- The baseline ----- # + + +def test_a_reform_leaves_its_baseline_and_the_baselines_caches_alone(): + system = new_system() + warm(system) + expected = read(system.parameters, INSTANT) + cached = cached_nodes(system.parameters) + at_instant = system.get_parameters_at_instant(INSTANT) + + class edits_the_tree(Reform): + def apply(self): + self.modify_parameters(modify_many_ways) + + def modify_many_ways(parameters): + parameters.benefit.amount.update(period=f"year:{INSTANT}:1", value=-1) + parameters.benefit.amount.values_list[-1].value = -2 + parameters.tax.brackets[1].rate.values_list[0].value = -3 + parameters.tax.brackets.pop(0) + parameters.allowance.children.pop("single") + parameters.benefit.add_child( + "new", ParameterNode(data={"leaf": {"2015-01-01": 1}}) + ) + return parameters + + edits_the_tree(system) + + assert_keeps_its_cached_nodes(cached) + assert system.get_parameters_at_instant(INSTANT) is at_instant + assert read(system.parameters.clone(), INSTANT) == expected + + +def test_replacing_metadata_and_list_values_stays_out_of_the_baseline(): + """A clone shares nested metadata and list values with what it was cloned + from, so a modifier replaces them. Editing one in place, which the deep copy + allowed, would reach the baseline: see ``modify_parameters``.""" + system = new_system() + system.parameters.add_child( + "sources", ParameterNode(data={"list": {"2015-01-01": ["salary"]}}) + ) + + def modify(parameters): + amount = parameters.benefit.amount + amount.metadata["reference"] = [*amount.metadata["reference"], "new"] + amount.metadata["new"] = True + entry = parameters.sources.list.values_list[0] + entry.value = [*entry.value, "pension"] + return parameters + + class replaces(Reform): + def apply(self): + self.modify_parameters(modify) + + reformed = replaces(system) + + assert reformed.parameters.sources.list(INSTANT) == ["salary", "pension"] + assert len(reformed.parameters.benefit.amount.metadata["reference"]) == 2 + assert system.parameters.sources.list(INSTANT) == ["salary"] + assert len(system.parameters.benefit.amount.metadata["reference"]) == 1 + assert "new" not in system.parameters.benefit.amount.metadata + + +def test_a_reform_of_a_system_with_no_parameters_can_supply_a_tree(): + system = new_system() + system.parameters = None + + class supplies_a_tree(Reform): + def apply(self): + self.modify_parameters( + lambda parameters: ParameterNode(data={"leaf": {"2015-01-01": 1}}) + ) + + reformed = supplies_a_tree(system) + + assert reformed.get_parameters_at_instant(instant("2016-01-01")).leaf == 1 diff --git a/tests/core/test_reform_modify_parameters_property.py b/tests/core/test_reform_modify_parameters_property.py new file mode 100644 index 000000000..98e89a8a8 --- /dev/null +++ b/tests/core/test_reform_modify_parameters_property.py @@ -0,0 +1,213 @@ +"""Properties of ``modify_parameters``. + +For any baseline read at any instants, and any reform whose ``apply()`` edits +the tree directly, reads it, calls modifiers ending in any way and uses the dict +form, on either path a reform runs: + +* Fresh reads: once ``modify_parameters`` has run last, the system, its tree and + each bracket read as a tree never read at all, given the same edits. That + holds for a modifier that raises as well. +* The two paths agree: an instance and ``apply_reform_set`` leave the same tree. +* The baseline stays put: an instance changes nothing in its baseline, its + values, its metadata or the nodes it has cached. +* Composition: two modifiers do what one modifier making both sets of edits does. +* Differential: for one modifier on a baseline never read, the tree matches what + the deep copy in ``Reform.modify_parameters`` used to leave. Only this case + can match: the deep copy threw away earlier steps, and read stale values. + +Two things are outside these properties, both on purpose. A direct edit to a tree +already read, with no ``modify_parameters`` after it, can read stale: #621 is +where that gets settled. A modifier that edits nested metadata or a list value in +place reaches the baseline, as ``modify_parameters`` documents: the edits here +replace top-level metadata keys only, and the tree has no list values. + +Examples are in ``test_reform_modify_parameters.py``. +""" + +import copy + +import pytest + +# The smoke job installs Core without the dev extra but collects every module. +pytest.importorskip("hypothesis") + +from hypothesis import given, settings # noqa: E402 +from hypothesis import strategies as st # noqa: E402 + +from policyengine_core.parameters import ParameterNode # noqa: E402 +from policyengine_core.reforms import Reform # noqa: E402 +from tests.fixtures.parameter_metadata_clones import INSTANTS, snapshot # noqa: E402 +from tests.fixtures.reform_modifiers import ( # noqa: E402 + ModifierFailed, + apply_steps, + assert_keeps_its_cached_nodes, + cached_nodes, + modifier, + new_system, + reads, + run_as_instance, + run_in_place, + system_reads, + warm, +) + +# Positions are taken modulo what the tree holds. +POSITIONS = st.integers(min_value=0, max_value=60) +YEARS = st.integers(2014, 2024) +NUMBERS = st.one_of(st.integers(0, 500), st.floats(0, 500, allow_nan=False)) + +EDITS = st.one_of( + st.tuples(st.just("update"), POSITIONS, YEARS, NUMBERS), + st.tuples(st.just("set value"), POSITIONS, POSITIONS, NUMBERS), + st.tuples(st.just("remove value"), POSITIONS, POSITIONS), + st.tuples(st.just("delete bracket"), POSITIONS, POSITIONS), + st.tuples(st.just("append bracket"), POSITIONS, NUMBERS, NUMBERS), + st.tuples(st.just("add child"), POSITIONS, NUMBERS), + st.tuples(st.just("remove child"), POSITIONS, POSITIONS), + st.tuples( + st.just("set key"), + POSITIONS, + st.sampled_from(["label", "unit", "added"]), + st.text(max_size=3), + ), + st.tuples(st.just("read"), POSITIONS), +) +EDIT_LISTS = st.lists(EDITS, max_size=4) +ENDINGS = st.sampled_from(["same tree", "deep copy", "clone", "no tree"]) + +MODIFY_STEPS = st.one_of( + st.tuples(st.just("modify"), EDIT_LISTS, ENDINGS), + st.tuples(st.just("dict"), YEARS, NUMBERS), +) +STEPS = st.lists( + st.one_of( + MODIFY_STEPS, + st.tuples(st.just("direct"), EDITS), + st.tuples(st.just("read"), POSITIONS), + ), + max_size=3, +) +PATHS = st.sampled_from([run_as_instance, run_in_place]) +WARMED = st.lists(st.sampled_from(INSTANTS), unique=True).map(tuple) + +_SETTINGS = dict(max_examples=200, deadline=None, derandomize=True) + + +def without_reads(steps): + """The same steps, with nothing in them that can fill a cache.""" + kept = [] + for step in steps: + if step[0] == "modify": + edits = [edit for edit in step[1] if edit[0] != "read"] + kept.append(("modify", edits, step[2])) + elif step[0] == "direct": + if step[1][0] != "read": + kept.append(step) + elif step[0] != "read": + kept.append(step) + return kept + + +def never_read(steps): + """A system given ``steps`` with no cache ever filled: the reference.""" + return run_in_place(new_system(), without_reads(steps)) + + +def warmed_system(instants): + system = new_system() + warm(system, instants) + return system + + +@settings(**_SETTINGS) +@given(instants=WARMED, steps=STEPS, last=MODIFY_STEPS, run=PATHS) +def test_reads_after_modify_parameters_match_a_tree_never_read( + instants, steps, last, run +): + steps = [*steps, last] + + reformed = run(warmed_system(instants), steps) + + assert system_reads(reformed) == system_reads(never_read(steps)) + + +@settings(**_SETTINGS) +@given(instants=WARMED, steps=STEPS, edits=EDIT_LISTS) +def test_a_modifier_that_raises_leaves_a_tree_that_reads_fresh(instants, steps, edits): + system = run_in_place(warmed_system(instants), steps) + + with pytest.raises(ModifierFailed): + system.modify_parameters(modifier(edits, "raises")) + + reference = never_read([*steps, ("modify", edits, "same tree")]) + assert system_reads(system) == system_reads(reference) + + +@settings(**_SETTINGS) +@given(instants=WARMED, steps=STEPS) +def test_the_two_paths_leave_the_same_tree(instants, steps): + instance = run_as_instance(warmed_system(instants), steps) + in_place = run_in_place(warmed_system(instants), steps) + + assert snapshot(instance.parameters) == snapshot(in_place.parameters) + + +@settings(**_SETTINGS) +@given(instants=WARMED, steps=STEPS, raises=st.booleans()) +def test_a_reform_instance_changes_nothing_in_its_baseline(instants, steps, raises): + system = warmed_system(instants) + before = snapshot(system.parameters) + cached = cached_nodes(system.parameters) + + if raises: + steps = [*steps, ("modify", [("set value", 0, 0, 1)], "raises")] + with pytest.raises(ModifierFailed): + run_as_instance(system, steps) + else: + run_as_instance(system, steps) + + assert snapshot(system.parameters) == before + assert_keeps_its_cached_nodes(cached) + assert reads(system.parameters) == reads(new_system().parameters) + + +@settings(**_SETTINGS) +@given(instants=WARMED, first=EDIT_LISTS, second=EDIT_LISTS, run=PATHS) +def test_two_modifiers_make_the_edits_of_one(instants, first, second, run): + apart = run( + warmed_system(instants), + [("modify", first, "same tree"), ("modify", second, "same tree")], + ) + together = run(warmed_system(instants), [("modify", first + second, "same tree")]) + + assert snapshot(apart.parameters) == snapshot(together.parameters) + assert system_reads(apart) == system_reads(together) + + +class DeepCopyingReform(Reform): + """A reform with ``modify_parameters`` as Core had it before #618.""" + + def modify_parameters(self, modifier_function): + baseline_parameters_copy = copy.deepcopy(self.baseline.parameters) + reform_parameters = modifier_function(baseline_parameters_copy) + if not isinstance(reform_parameters, ParameterNode): + return ValueError("modifier_function must return a ParameterNode") + self.parameters = reform_parameters + self._parameters_at_instant_cache = {} + + +@settings(**_SETTINGS) +@given(edits=EDIT_LISTS, ending=st.sampled_from(["same tree", "deep copy", "clone"])) +def test_one_modifier_matches_the_deep_copy_on_a_baseline_never_read(edits, ending): + edits = [edit for edit in edits if edit[0] != "read"] + steps = [("modify", edits, ending)] + + class deep_copying(DeepCopyingReform): + def apply(self): + apply_steps(self, steps) + + reference = deep_copying(new_system()) + reformed = run_as_instance(new_system(), steps) + + assert snapshot(reformed.parameters) == snapshot(reference.parameters) + assert system_reads(reformed) == system_reads(reference) diff --git a/tests/core/test_reforms.py b/tests/core/test_reforms.py index 711888b50..2cf55fa26 100644 --- a/tests/core/test_reforms.py +++ b/tests/core/test_reforms.py @@ -282,7 +282,6 @@ def modify_parameters(reference_parameters): }, ) reference_parameters.children["new_node"] = reform_parameters_subtree - reference_parameters._at_instant_cache.clear() return reference_parameters class test_modify_parameters(Reform): diff --git a/tests/core/test_tracing_parameter_isolation.py b/tests/core/test_tracing_parameter_isolation.py index 8c4f5e95a..ef3d717a4 100644 --- a/tests/core/test_tracing_parameter_isolation.py +++ b/tests/core/test_tracing_parameter_isolation.py @@ -16,7 +16,7 @@ - ``ParameterNode.clone`` copies ``trace`` and ``tracer``, so every clone of the system traced as well; - copying a ``TracingParameterNodeAtInstant`` (``copy``, ``deepcopy`` or - ``pickle``, as ``Reform.modify_parameters`` does when it deep-copies the + ``pickle``, as ``Reform.modify_parameters`` did when it deep-copied the tree) recursed without end, because ``__getattr__`` read ``parameter_node_at_instant`` from an instance the copy protocol had not filled in yet. @@ -78,10 +78,11 @@ def test_traced_calculation_leaves_the_system_untraced(isolated_tax_benefit_syst assert_untraced(system.clone().parameters) -def test_reform_after_a_traced_calculation_can_copy_the_parameters( +def test_reform_after_a_traced_calculation_reads_its_own_parameters( isolated_tax_benefit_system, ): - """``Reform.modify_parameters`` deep-copies the baseline's parameter tree.""" + """A reform built after a traced calculation reads its own parameters and + leaves the system's unchanged.""" system = isolated_tax_benefit_system simulation = build_simulation(system, trace=True) simulation.calculate("income_tax", JANUARY) diff --git a/tests/fixtures/reform_modifiers.py b/tests/fixtures/reform_modifiers.py new file mode 100644 index 000000000..1ff5c79b4 --- /dev/null +++ b/tests/fixtures/reform_modifiers.py @@ -0,0 +1,247 @@ +"""Systems, modifiers and reads for the ``modify_parameters`` tests. + +``new_system`` gives a tax-benefit system holding the tree from +``parameter_metadata_clones.tree_data``, which has nodes, parameters, a scale +and brackets. Its brackets have no ``parent``, as loaded brackets don't. + +``read`` turns everything a tree returns at an instant into plain data, through +each cache a read can come from: the root, the nodes below it and the brackets. +``warm`` fills those caches. + +``make_edit`` makes one edit to a tree. It adds three kinds to +``parameter_metadata_clones.apply_edit``, none of which clears a cache: deleting +a bracket, appending one, and adding or removing a child. These with assigning +and removing a value are what ``Parameter.update`` doesn't see. + +``make_edit`` also takes ``("read", position)``, which reads the tree at one of +``INSTANTS``: the only edit that fills a cache. + +``modifier`` turns edits and an ending into a modifier function. The ending says +what it does once it has made its edits: hand back the tree it got, or a copy of +it, or something that isn't a tree, or raise. + +``run_as_instance`` and ``run_in_place`` are the two ways a reform runs. See +#618. +""" + +import copy + +from policyengine_core.country_template import CountryTaxBenefitSystem +from policyengine_core.parameters import ( + ParameterNode, + ParameterScale, + ParameterScaleBracket, +) +from policyengine_core.parameters.parameter_node_at_instant import ( + ParameterNodeAtInstant, +) +from policyengine_core.reforms import Reform +from policyengine_core.taxscales import TaxScaleLike +from tests.fixtures.parameter_metadata_clones import ( + INSTANTS, + apply_edit, + build, + parameters as parameters_of, +) + +STRUCTURAL_EDITS = ("delete bracket", "append bracket", "add child", "remove child") + + +class ModifierFailed(Exception): + """Raised by a modifier whose ending is ``"raises"``.""" + + +def new_system(): + system = CountryTaxBenefitSystem() + system.parameters = build() + system._parameters_at_instant_cache = {} + return system + + +def _read_at_instant(item): + if isinstance(item, ParameterNodeAtInstant): + return {key: _read_at_instant(item[key]) for key in item._children} + if isinstance(item, TaxScaleLike): + return ( + type(item).__name__, + [float(threshold) for threshold in item.thresholds], + [float(rate) for rate in getattr(item, "rates", [])], + [float(amount) for amount in getattr(item, "amounts", [])], + ) + return item + + +def _brackets(root): + return [ + bracket + for item in root.get_descendants() + if isinstance(item, ParameterScale) + for bracket in item.brackets + ] + + +def read(root, instant_str): + """What ``root`` and each bracket under it return at an instant.""" + return ( + _read_at_instant(root(instant_str)), + [_read_at_instant(bracket(instant_str)) for bracket in _brackets(root)], + ) + + +def reads(root): + return [read(root, instant_str) for instant_str in INSTANTS] + + +def system_reads(system): + """Every instant as a system returns it, then as its tree returns it.""" + from_system = [ + _read_at_instant(system.get_parameters_at_instant(instant_str)) + for instant_str in INSTANTS + ] + return from_system, reads(system.parameters) + + +def warm(system, instant_strs=INSTANTS): + """Fill the caches of ``system`` and its tree at ``instant_strs``.""" + for instant_str in instant_strs: + system.get_parameters_at_instant(instant_str) + read(system.parameters, instant_str) + + +def _nodes(root): + return [root] + [ + item + for item in root.get_descendants() + if isinstance(item, ParameterNode) + and not isinstance(item, ParameterScaleBracket) + ] + + +def _scales(root): + return [item for item in root.get_descendants() if isinstance(item, ParameterScale)] + + +def make_edit(root, edit): + """Make ``edit`` on the tree under ``root``. + + Positions are taken modulo what's in the tree, so any position names + something in any tree. An edit with nothing to act on does nothing. + """ + kind = edit[0] + if kind == "read": + read(root, INSTANTS[edit[1] % len(INSTANTS)]) + return + if kind not in STRUCTURAL_EDITS: + if kind in ("update", "set value", "remove value") and not parameters_of(root): + return + apply_edit(root, edit) + return + if kind == "add child": + _, position, value = edit + node = _nodes(root)[position % len(_nodes(root))] + name = f"added_{len(node.children)}" + while name in node.children: + name += "_" + node.add_child( + name, + ParameterNode(data={"leaf": {"2015-01-01": value}}), + ) + elif kind == "remove child": + node = _nodes(root)[edit[1] % len(_nodes(root))] + if node.children: + name = sorted(node.children)[edit[2] % len(node.children)] + del node.children[name] + delattr(node, name) + elif _scales(root): + scale = _scales(root)[edit[1] % len(_scales(root))] + if kind == "delete bracket": + if scale.brackets: + del scale.brackets[edit[2] % len(scale.brackets)] + else: + _, _, threshold, rate = edit + scale.brackets.append( + ParameterScaleBracket( + "new_bracket", + data={ + "threshold": {"2015-01-01": {"value": threshold}}, + "rate": {"2015-01-01": {"value": rate}}, + }, + ) + ) + + +def modifier(edits, ending="same tree"): + """A modifier function that makes ``edits``, then does as ``ending`` says.""" + + def modify(parameters): + if ending == "deep copy": + parameters = copy.deepcopy(parameters) + elif ending == "clone": + parameters = parameters.clone() + for edit in edits: + make_edit(parameters, edit) + if ending == "no tree": + return None + if ending == "raises": + raise ModifierFailed + return parameters + + return modify + + +def apply_steps(system, steps): + """Make each step of a reform's ``apply()`` on ``system``. + + A step is an edit made straight on the system's tree (``("direct", edit)``), + a read of it (``("read", position)``), a modifier function (``("modify", + edits, ending)``), or the dict form of ``modify_parameters`` (``("dict", + year, value)``), which does nothing once an edit has removed its parameter. + """ + for step in steps: + if step[0] == "direct": + make_edit(system.parameters, step[1]) + elif step[0] == "read": + make_edit(system.parameters, step) + elif step[0] == "modify": + system.modify_parameters(modifier(step[1], step[2])) + else: + _, year, value = step + benefit = system.parameters.children.get("benefit") + changes = {} + if benefit is not None and "amount" in benefit.children: + changes["benefit.amount"] = {f"year:{year}:1": value} + system.modify_parameters(changes) + + +def reform_with(steps): + class reform(Reform): + def apply(self): + apply_steps(self, steps) + + return reform + + +def run_as_instance(system, steps): + """Run a reform as ``MyReform(baseline)``; return the reform.""" + return reform_with(steps)(system) + + +def run_in_place(system, steps): + """Run a reform with ``system.apply_reform_set``; return the system.""" + system.apply_reform_set(reform_with(steps)) + return system + + +def cached_nodes(root): + """Each node of a tree with what it has cached, to check nothing replaced it.""" + return { + id(item): (item, dict(item._at_instant_cache)) + for item in _nodes(root) + _brackets(root) + } + + +def assert_keeps_its_cached_nodes(cached): + for item, nodes in cached.values(): + assert item._at_instant_cache.keys() == nodes.keys() + for instant_str, node in nodes.items(): + assert item._at_instant_cache[instant_str] is node From c69bd859dfbf10c99fc3ced770ab170fecc63cc4 Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Mon, 2 Nov 2026 04:00:59 -0500 Subject: [PATCH 2/2] Pin what the reviews of modify_parameters asked about Add tests for the system cache being rebound after the modifier, a read through the system inside a modifier, the return value, a walk that stops at the nodes below, and edits made before a modifier raises. Warm the tree inside apply() so the instance path exercises the walk as well. Co-Authored-By: Claude Opus 5.5 --- .../taxbenefitsystems/tax_benefit_system.py | 4 +- tests/core/test_reform_modify_parameters.py | 115 +++++++++++++++--- tests/fixtures/reform_modifiers.py | 8 +- 3 files changed, 107 insertions(+), 20 deletions(-) diff --git a/policyengine_core/taxbenefitsystems/tax_benefit_system.py b/policyengine_core/taxbenefitsystems/tax_benefit_system.py index 3dbf45447..5c0952970 100644 --- a/policyengine_core/taxbenefitsystems/tax_benefit_system.py +++ b/policyengine_core/taxbenefitsystems/tax_benefit_system.py @@ -796,7 +796,7 @@ 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. @@ -810,7 +810,7 @@ def modify_parameters( 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. + 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. diff --git a/tests/core/test_reform_modify_parameters.py b/tests/core/test_reform_modify_parameters.py index e471ddb5a..9951c6e26 100644 --- a/tests/core/test_reform_modify_parameters.py +++ b/tests/core/test_reform_modify_parameters.py @@ -156,15 +156,20 @@ def modify(parameters): edit(parameters) return parameters - class edited(Reform): - def apply(self): - self.modify_parameters(modify) + def reform(warm_first): + class edited(Reform): + def apply(self): + if warm_first: + warm(self) + self.modify_parameters(modify) + + return edited warmed = new_system() warm(warmed) - reformed = run(edited, warmed) + reformed = run(reform(True), warmed) - assert system_reads(reformed) == system_reads(run(edited, new_system())) + assert system_reads(reformed) == system_reads(run(reform(False), new_system())) def test_dict_form_clears_what_an_update_under_a_bracket_leaves_behind(): @@ -172,10 +177,10 @@ def test_dict_form_clears_what_an_update_under_a_bracket_leaves_behind(): system = new_system() warm(system) - system.modify_parameters({"tax[1].rate": {"year:2017:1": 0.9}}) + system.modify_parameters({"tax[1].rate": {"year:2016:1": 0.9}}) - assert system.get_parameters_at_instant(INSTANT).tax.rates == [0.1, 0.9] - assert system.parameters(INSTANT).tax.rates == [0.1, 0.9] + assert system.get_parameters_at_instant("2016-01-01").tax.rates == [0.1, 0.9] + assert system.parameters("2016-01-01").tax.rates == [0.1, 0.9] @BOTH_PATHS @@ -190,6 +195,7 @@ def modify(parameters): class replaces_the_tree(Reform): def apply(self): + warm(self) self.modify_parameters(modify) system = new_system() @@ -234,19 +240,79 @@ def modify(parameters): assert system.parameters("2020-01-01").benefit.amount == -1 -def test_clear_descendant_caches_reaches_every_node_below_and_none_above(): +@BOTH_PATHS +def test_edits_made_before_a_modifier_raises_stay(run): + """``apply()`` may catch the exception and carry on.""" + + def modify(parameters): + parameters.benefit.amount.values_list[0].value = -1 + raise ModifierFailed + + class catches_it(Reform): + def apply(self): + warm(self) + with pytest.raises(ModifierFailed): + self.modify_parameters(modify) + + reformed = run(catches_it, new_system()) + + assert reformed.get_parameters_at_instant("2020-01-01").benefit.amount == -1 + + +@BOTH_PATHS +def test_a_read_through_the_system_inside_a_modifier_leaves_nothing_stale(run): + class reads_then_edits(Reform): + def apply(self): + def modify(parameters): + self.get_parameters_at_instant("2020-01-01") + parameters.benefit.amount.values_list[0].value = -1 + return parameters + + self.modify_parameters(modify) + + reformed = run(reads_then_edits, new_system()) + + assert reformed.get_parameters_at_instant("2020-01-01").benefit.amount == -1 + + +def test_the_system_gets_a_new_cache_and_leaves_the_old_one_alone(): + """policyengine-us's SPM policies share this dict, so emptying it in place + would reach the others.""" system = new_system() warm(system) + shared = system._parameters_at_instant_cache + instants = set(shared) + + system.modify_parameters(lambda parameters: parameters) + + assert instants + assert system._parameters_at_instant_cache is not shared + assert set(shared) == instants + + +def test_clear_descendant_caches_reaches_every_node_below_and_no_other(): + system = new_system() parameters = system.parameters + parameters.benefit.add_child( + "inner", ParameterNode(data={"node": {"leaf": {"2015-01-01": 1}}}) + ) + warm(system) + below = [ + parameters.benefit, + parameters.benefit.inner, + parameters.benefit.inner.node, + ] + nodes = [node for node, _ in cached_nodes(parameters).values()] + assert all(node._at_instant_cache for node in nodes) parameters.benefit.clear_descendant_caches() - assert not parameters.benefit._at_instant_cache - assert parameters._at_instant_cache + + for node in nodes: + assert bool(node._at_instant_cache) != any(node is item for item in below) parameters.clear_descendant_caches() - assert not any( - node._at_instant_cache for node, _ in cached_nodes(parameters).values() - ) + + assert not any(node._at_instant_cache for node in nodes) # ----- Redundant copy, and the changes it dropped ----- # @@ -338,10 +404,11 @@ def test_a_modifier_that_returns_no_tree_keeps_its_edits(run): results = [] def modify(parameters): - parameters.benefit.amount.update(period="year:2020:1", value=-1) + parameters.benefit.amount.values_list[0].value = -1 class returns_nothing(Reform): def apply(self): + warm(self) results.append(self.modify_parameters(modify)) system = new_system() @@ -353,6 +420,24 @@ def apply(self): assert reformed.get_parameters_at_instant("2020-01-01").benefit.amount == -1 +@BOTH_PATHS +@pytest.mark.parametrize( + "change", + [lambda parameters: parameters, {"benefit.amount": {"year:2020:1": -1}}], + ids=["function", "dict"], +) +def test_modify_parameters_returns_the_system_it_ran_on(run, change): + results = [] + + class records(Reform): + def apply(self): + results.append((self.modify_parameters(change), self)) + + run(records, new_system()) + + assert results[0][0] is results[0][1] + + def test_a_reform_instance_takes_the_dict_form(): class dict_form(Reform): def apply(self): diff --git a/tests/fixtures/reform_modifiers.py b/tests/fixtures/reform_modifiers.py index 1ff5c79b4..c5b3a649b 100644 --- a/tests/fixtures/reform_modifiers.py +++ b/tests/fixtures/reform_modifiers.py @@ -193,14 +193,16 @@ def apply_steps(system, steps): """Make each step of a reform's ``apply()`` on ``system``. A step is an edit made straight on the system's tree (``("direct", edit)``), - a read of it (``("read", position)``), a modifier function (``("modify", - edits, ending)``), or the dict form of ``modify_parameters`` (``("dict", - year, value)``), which does nothing once an edit has removed its parameter. + a read through the system and its tree (``("read", position)``), a modifier + function (``("modify", edits, ending)``), or the dict form of + ``modify_parameters`` (``("dict", year, value)``), which does nothing once an + edit has removed its parameter. """ for step in steps: if step[0] == "direct": make_edit(system.parameters, step[1]) elif step[0] == "read": + system.get_parameters_at_instant(INSTANTS[step[1] % len(INSTANTS)]) make_edit(system.parameters, step) elif step[0] == "modify": system.modify_parameters(modifier(step[1], step[2]))