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..5c0952970 100644 --- a/policyengine_core/taxbenefitsystems/tax_benefit_system.py +++ b/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): 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..9951c6e26 --- /dev/null +++ b/tests/core/test_reform_modify_parameters.py @@ -0,0 +1,527 @@ +"""``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 + + 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(reform(True), warmed) + + assert system_reads(reformed) == system_reads(run(reform(False), 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:2016: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 +@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): + warm(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 + + +@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() + + 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 nodes) + + +# ----- 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.values_list[0].value = -1 + + class returns_nothing(Reform): + def apply(self): + warm(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 + + +@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): + 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..c5b3a649b --- /dev/null +++ b/tests/fixtures/reform_modifiers.py @@ -0,0 +1,249 @@ +"""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 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])) + 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