From b6c68eb363fd746b60e2a7c8dacfefcbd622840e Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Fri, 2 Oct 2026 09:17:13 -0400 Subject: [PATCH] Let simulations be deep-copied and unpickled copy and pickle rebuild an object by creating an empty instance and probing it for __setstate__ before restoring its __dict__. Population.__getattr__ answered that probe through the projector lookup, which reads self.entity; on the empty instance that read re-entered __getattr__ until RecursionError. So copy.deepcopy(simulation), pickle.loads(pickle.dumps(simulation)) and copy.copy(population) failed on every simulation, on Python 3.11 to 3.14. The same path had three more defects: - VectorialParameterNodeAtInstant and TracingParameterNodeAtInstant recursed the same way through the attribute they forward to. - The vectorial node forwarded __deepcopy__ to its numpy vector, so a deep copy came back as a bare recarray. These nodes are cached on the parameter nodes a tax-benefit system keeps. - numpy's __reduce__ rebuilt an EnumArray without possible_values, so an unpickled enum array could be neither decoded nor compared with an enum item. EnumArray now pickles its enum by name and restores it when the process can find it, otherwise None. Tests: examples for each defect, plus a Hypothesis property (in its own importorskip module, for the smoke job) that a deep copy or pickle round trip calculates what a freshly built simulation does and that writes to the copy never reach the original. Co-Authored-By: Claude Opus 5.5 --- ...-simulation-copy-pickle-recursion.fixed.md | 1 + policyengine_core/enums/enum_array.py | 38 +- .../vectorial_parameter_node_at_instant.py | 8 + policyengine_core/populations/population.py | 8 + .../tracing_parameter_node_at_instant.py | 4 + tests/core/test_simulation_copy_pickle.py | 341 ++++++++++++++++++ .../test_simulation_copy_pickle_property.py | 99 +++++ tests/fixtures/simulation_copy.py | 81 +++++ 8 files changed, 579 insertions(+), 1 deletion(-) create mode 100644 changelog.d/fix-simulation-copy-pickle-recursion.fixed.md create mode 100644 tests/core/test_simulation_copy_pickle.py create mode 100644 tests/core/test_simulation_copy_pickle_property.py create mode 100644 tests/fixtures/simulation_copy.py diff --git a/changelog.d/fix-simulation-copy-pickle-recursion.fixed.md b/changelog.d/fix-simulation-copy-pickle-recursion.fixed.md new file mode 100644 index 00000000..1332e922 --- /dev/null +++ b/changelog.d/fix-simulation-copy-pickle-recursion.fixed.md @@ -0,0 +1 @@ +Stop `copy.deepcopy(simulation)`, `copy.copy(population)` and unpickling a simulation in the process that pickled it from raising `RecursionError`, keep vectorial parameter nodes as nodes when deep-copied, and keep an `EnumArray`'s `possible_values` through pickling. diff --git a/policyengine_core/enums/enum_array.py b/policyengine_core/enums/enum_array.py index c518add8..36be3f03 100644 --- a/policyengine_core/enums/enum_array.py +++ b/policyengine_core/enums/enum_array.py @@ -1,7 +1,9 @@ from __future__ import annotations +import importlib +import operator import typing -from typing import Any, NoReturn, Optional, Type +from typing import Any, NoReturn, Optional, Tuple, Type import numpy @@ -9,6 +11,31 @@ from policyengine_core.enums import Enum +def _restore_enum_array( + array: numpy.ndarray, enum_name: Optional[Tuple[str, str]] +) -> EnumArray: + """Rebuild a pickled EnumArray, with its enum if this process can find it. + + ``enum_name`` is the enum's module and qualified name, or ``None`` for an + array that had no enum. + + Tax-benefit systems load variable files under module names that exist only + in the process that loaded them, so an enum defined in one cannot be found + from another process. The array then comes back with ``possible_values`` + unset (``None``), where pickling the enum by reference would fail to + unpickle the array at all. + """ + possible_values = None + if enum_name is not None: + module_name, qualified_name = enum_name + try: + module = importlib.import_module(module_name) + possible_values = operator.attrgetter(qualified_name)(module) + except (ImportError, AttributeError): + pass + return EnumArray(array, possible_values) + + class EnumArray(numpy.ndarray): """ Numpy array subclass representing an array of enum items. @@ -35,6 +62,15 @@ def __array_finalize__(self, obj: Optional[numpy.int_]) -> None: self.possible_values = getattr(obj, "possible_values", None) + def __reduce__(self) -> tuple: + # ndarray's own ``__reduce__`` rebuilds the array without + # ``possible_values``, so an unpickled EnumArray could be neither + # decoded nor compared with an enum item. The enum travels by name + # rather than by reference; see ``_restore_enum_array``. + enum = self.possible_values + name = None if enum is None else (enum.__module__, enum.__qualname__) + return _restore_enum_array, (self.view(numpy.ndarray), name) + def __eq__(self, other: Any) -> bool: # When comparing to an item of self.possible_values, use the item index # to speed up the comparison. diff --git a/policyengine_core/parameters/vectorial_parameter_node_at_instant.py b/policyengine_core/parameters/vectorial_parameter_node_at_instant.py index fcc562ef..81f1a9d3 100644 --- a/policyengine_core/parameters/vectorial_parameter_node_at_instant.py +++ b/policyengine_core/parameters/vectorial_parameter_node_at_instant.py @@ -209,6 +209,14 @@ def __init__(self, name: str, vector: ArrayLike, instant_str: str): self._instant_str = instant_str def __getattr__(self, attribute: str) -> Any: + # ``vector`` is missing while copy or pickle rebuilds a node, and + # looking it up would recurse here. ``copy.deepcopy`` looks + # ``__deepcopy__`` up on the instance, and the vector's would copy the + # vector alone and hand back a bare ``numpy.recarray``. + if attribute in ("vector", "__deepcopy__"): + raise AttributeError( + f"{type(self).__name__!s} has no attribute {attribute!r}" + ) result = getattr(self.vector, attribute) if isinstance(result, numpy.recarray): return VectorialParameterNodeAtInstant(result) diff --git a/policyengine_core/populations/population.py b/policyengine_core/populations/population.py index d2e8df1d..d391b810 100644 --- a/policyengine_core/populations/population.py +++ b/policyengine_core/populations/population.py @@ -53,6 +53,14 @@ def filled_array(self, value: Any, dtype: Any = None) -> numpy.ndarray: return numpy.full(self.count, value, dtype) def __getattr__(self, attribute: str) -> Any: + # The shortcut lookup reads ``self.entity`` and ``self.simulation``. + # They are missing while copy or pickle rebuilds a population (both + # probe the new, empty instance for ``__setstate__`` before restoring + # its ``__dict__``), and looking them up would recurse here. + if attribute in ("entity", "simulation"): + raise AttributeError( + f"{type(self).__name__!s} has no attribute {attribute!r}" + ) projector = projectors.get_projector_from_shortcut(self, attribute) if not projector: raise AttributeError( diff --git a/policyengine_core/tracers/tracing_parameter_node_at_instant.py b/policyengine_core/tracers/tracing_parameter_node_at_instant.py index d4bf1746..5d3fde49 100644 --- a/policyengine_core/tracers/tracing_parameter_node_at_instant.py +++ b/policyengine_core/tracers/tracing_parameter_node_at_instant.py @@ -36,6 +36,10 @@ def __getattr__( self, key: str, ) -> Union[TracingParameterNodeAtInstant, Child]: + # ``parameter_node_at_instant`` is missing while copy or pickle + # rebuilds a wrapper, and looking it up would recurse here. + if key == "parameter_node_at_instant": + raise AttributeError(f"{type(self).__name__!s} has no attribute {key!r}") child = getattr(self.parameter_node_at_instant, key) return self.get_traced_child(child, key) diff --git a/tests/core/test_simulation_copy_pickle.py b/tests/core/test_simulation_copy_pickle.py new file mode 100644 index 00000000..89fc212a --- /dev/null +++ b/tests/core/test_simulation_copy_pickle.py @@ -0,0 +1,341 @@ +"""Simulations survive ``copy.deepcopy`` and a pickle round trip. + +``copy`` and ``pickle`` rebuild an object by creating an empty instance and +then probing it for ``__setstate__`` before its ``__dict__`` is restored. A +``__getattr__`` that reads one of the instance's own attributes recurses on +that probe until the interpreter raises ``RecursionError``. +``Population.__getattr__`` did this through its projector lookup (it reads +``self.entity``), so ``copy.deepcopy(simulation)``, +``pickle.loads(pickle.dumps(simulation))`` and even ``copy.copy(population)`` +failed on every simulation. The parameter wrappers that forward attribute +lookups (vectorial and tracing nodes) had the same recursion, and the +vectorial node also forwarded ``__deepcopy__`` to its numpy vector, so a deep +copy came back as a bare ``numpy.recarray``. Behind the recursion, a pickled +``EnumArray`` lost its ``possible_values``. + +Pickles are for the process that wrote them: a tax-benefit system loads +variable files under module names no other process has. +""" + +from __future__ import annotations + +import copy +import io +import pickle + +import numpy as np +import pytest + +from policyengine_core.country_template import CountryTaxBenefitSystem +from policyengine_core.enums import Enum, EnumArray +from policyengine_core.parameters import VectorialParameterNodeAtInstant +from policyengine_core.populations import GroupPopulation, Population +from policyengine_core.tools import assert_near +from policyengine_core.tracers import FullTracer, TracingParameterNodeAtInstant +from tests.fixtures.simulation_copy import ( + COPIERS, + FEB, + JAN, + DoubleIncomeTaxRate, + build_simulation, + rate_node, +) + +# Population: lookups on an instance that copy or pickle has not filled yet. + + +@pytest.mark.parametrize("population_class", [Population, GroupPopulation]) +def test_lookups_on_an_unfilled_population_raise_attribute_error(population_class): + unfilled = population_class.__new__(population_class) + + assert getattr(unfilled, "__setstate__", None) is None + assert not hasattr(unfilled, "household") + with pytest.raises(AttributeError): + unfilled.entity + with pytest.raises(AttributeError): + unfilled.simulation + + +def test_copy_protocol_lookups_on_a_population_find_nothing(): + simulation = build_simulation() + + assert getattr(simulation.persons, "__deepcopy__", None) is None + assert getattr(simulation.household, "__setstate__", None) is None + + +def test_population_projector_shortcuts_still_resolve(): + simulation = build_simulation() + + assert_near( + simulation.persons.household("rent", JAN), simulation.household("rent", JAN) + ) + with pytest.raises(AttributeError, match="not a known attribute"): + simulation.persons.not_an_entity + + +def test_shallow_copy_of_a_population_shares_its_holders(): + simulation = build_simulation() + + population = copy.copy(simulation.persons) + + assert population.entity is simulation.persons.entity + assert population._holders is simulation.persons._holders + + +# Simulations. + + +@pytest.mark.parametrize("copier", COPIERS.values(), ids=COPIERS.keys()) +def test_copied_simulation_calculates_what_the_original_does(copier): + simulation = build_simulation(salaries=(1000, 2500)) + expected = simulation.calculate("income_tax", JAN) + + copied = copier(simulation) + + assert_near(copied.calculate("income_tax", JAN), expected) + assert_near( + copied.calculate("disposable_income", JAN), + simulation.calculate("disposable_income", JAN), + ) + assert copied.persons is not simulation.persons + assert copied.persons.simulation is copied + assert copied.household.simulation is copied + assert copied.household.members is copied.persons + assert copied.populations["person"] is copied.persons + + +@pytest.mark.parametrize("copier", COPIERS.values(), ids=COPIERS.keys()) +def test_inputs_set_on_a_copy_do_not_reach_the_original(copier): + simulation = build_simulation() + simulation.calculate("income_tax", JAN) + + copied = copier(simulation) + copied.set_input("salary", FEB, [2000]) + + assert_near(copied.calculate("income_tax", FEB), [300]) + assert simulation.persons.get_holder("salary").get_array(FEB) is None + assert simulation.persons.get_holder("income_tax").get_array(FEB) is None + + +@pytest.mark.parametrize("copier", COPIERS.values(), ids=COPIERS.keys()) +def test_copied_simulation_keeps_its_branches(copier): + simulation = build_simulation() + simulation.get_branch("child").set_input("salary", FEB, [2000]) + + copied = copier(simulation) + + child = copied.branches["child"] + assert child.parent_branch is copied + assert child is not simulation.branches["child"] + assert_near(child.calculate("income_tax", FEB), [300]) + + +@pytest.mark.parametrize("copier", COPIERS.values(), ids=COPIERS.keys()) +def test_copied_reform_simulation_keeps_the_reform(copier): + reformed = DoubleIncomeTaxRate(CountryTaxBenefitSystem()) + simulation = build_simulation(tax_benefit_system=reformed) + expected = simulation.calculate("income_tax", JAN) + + copied = copier(simulation) + + assert isinstance(copied.tax_benefit_system, DoubleIncomeTaxRate) + assert_near(copied.calculate("income_tax", JAN), expected) + assert_near(expected, [300]) + + +@pytest.mark.parametrize("copier", COPIERS.values(), ids=COPIERS.keys()) +def test_copied_traced_simulation_still_traces(copier): + simulation = build_simulation() + simulation.trace = True + simulation.calculate("income_tax", JAN) + + copied = copier(simulation) + copied.calculate("disposable_income", JAN) + + assert copied.tracer is not simulation.tracer + assert "disposable_income<2025-01, (default)>" in copied.tracer.get_flat_trace() + + +@pytest.mark.parametrize("copier", COPIERS.values(), ids=COPIERS.keys()) +def test_copied_simulation_keeps_cached_vectorial_parameter_nodes(copier): + tax_benefit_system = CountryTaxBenefitSystem() + tax_benefit_system.parameters.add_child("rate", rate_node()) + simulation = build_simulation(tax_benefit_system=tax_benefit_system) + statuses = np.array(["couple", "single"]) + parameters = simulation.tax_benefit_system.get_parameters_at_instant("2015-01-01") + assert_near(parameters.rate[statuses].owner, [500, 100]) # caches the node + + copied = copier(simulation) + + rate = copied.tax_benefit_system.get_parameters_at_instant("2015-01-01").rate + assert isinstance(rate._vectorial_node, VectorialParameterNodeAtInstant) + assert_near(rate[statuses].tenant, [700, 300]) + + +@pytest.mark.parametrize("copier", COPIERS.values(), ids=COPIERS.keys()) +def test_copied_simulation_keeps_enum_inputs(copier): + simulation = build_simulation( + salaries=(1000, 1000), + households=[[0], [1]], + occupancy=["owner", "free_lodger"], + ) + + copied = copier(simulation) + + status = copied.household("housing_occupancy_status", JAN) + assert list(status.decode_to_str()) == ["owner", "free_lodger"] + # The formula reads the enum from the stored array's ``possible_values``. + housing_tax = copied.calculate("housing_tax", "2025") + assert_near(housing_tax, simulation.calculate("housing_tax", "2025")) + assert housing_tax[0] > 0 + assert housing_tax[1] == 0 + + +# Enum arrays. + + +@pytest.mark.parametrize("protocol", range(pickle.HIGHEST_PROTOCOL + 1)) +def test_pickled_enum_array_keeps_its_possible_values(protocol): + status = build_simulation(occupancy=["free_lodger"]).household( + "housing_occupancy_status", JAN + ) + + restored = pickle.loads(pickle.dumps(status, protocol=protocol)) + + assert type(restored) is EnumArray + assert restored.possible_values is status.possible_values + assert list(restored.decode_to_str()) == ["free_lodger"] + assert (restored == status.possible_values.free_lodger).all() + assert restored.dtype == status.dtype + + +def test_pickled_enum_array_without_possible_values_stays_without(): + restored = pickle.loads(pickle.dumps(EnumArray(np.array([0, 1])))) + + assert type(restored) is EnumArray + assert restored.possible_values is None + assert restored.tolist() == [0, 1] + + +@pytest.mark.parametrize( + "module_name, qualified_name", + [ + ("module_only_another_process_loaded", "HousingOccupancyStatus"), + ("policyengine_core.enums", "NoSuchEnum"), + ], +) +def test_enum_array_whose_enum_cannot_be_found_unpickles_without_it( + module_name, qualified_name +): + """An enum from a variable file has a module name only its process knows.""" + + class Unreachable(Enum): + a = "a" + b = "b" + + Unreachable.__module__ = module_name + Unreachable.__qualname__ = qualified_name + array = EnumArray(np.array([1, 0, 1]), Unreachable) + + restored = pickle.loads(pickle.dumps(array)) + + assert type(restored) is EnumArray + assert restored.possible_values is None + assert restored.tolist() == [1, 0, 1] + + +def test_enum_array_pickled_by_an_earlier_release_still_loads(): + """Earlier releases pickled an EnumArray as numpy pickles any ndarray.""" + status = build_simulation(occupancy=["owner"]).household( + "housing_occupancy_status", JAN + ) + + class EarlierPickler(pickle.Pickler): + def reducer_override(self, obj): + if isinstance(obj, EnumArray): + return np.ndarray.__reduce__(obj) + return NotImplemented + + buffer = io.BytesIO() + EarlierPickler(buffer).dump(status) + + restored = pickle.loads(buffer.getvalue()) + + assert type(restored) is EnumArray + assert restored.tolist() == status.tolist() + + +def test_copied_enum_array_keeps_its_possible_values(): + status = build_simulation(occupancy=["owner"]).household( + "housing_occupancy_status", JAN + ) + + for copied in (copy.copy(status), copy.deepcopy(status)): + assert type(copied) is EnumArray + assert copied.possible_values is status.possible_values + assert list(copied.decode_to_str()) == ["owner"] + assert not np.shares_memory(copied, status) + + +# Parameter wrappers that forward attribute lookups. + + +@pytest.mark.parametrize("copier", COPIERS.values(), ids=COPIERS.keys()) +def test_copied_vectorial_parameter_node_is_a_vectorial_node(copier): + node = rate_node()("2015-01-01")[np.array(["single", "couple", "single"])] + + copied = copier(node) + + assert type(copied) is VectorialParameterNodeAtInstant + assert copied._name == node._name + assert copied._instant_str == node._instant_str + assert_near(copied.owner, [100, 500, 100]) + assert_near(copied[np.array(["owner", "tenant", "tenant"])], [100, 700, 300]) + assert not np.shares_memory(copied.vector, node.vector) + + +def test_vectorial_parameter_node_still_forwards_to_its_vector(): + node = rate_node()("2015-01-01")[np.array(["single", "couple"])] + + assert set(node.dtype.names) == {"owner", "tenant"} + assert node.shape == (2,) + + +@pytest.mark.parametrize("copier", COPIERS.values(), ids=COPIERS.keys()) +def test_copied_tracing_parameter_node_still_traces(copier): + tracer = FullTracer() + tracer.record_calculation_start("v", "2015-01") + node = TracingParameterNodeAtInstant( + rate_node()("2015-01-01"), tracer, branch_name="default" + ) + + copied = copier(node) + + assert type(copied) is TracingParameterNodeAtInstant + assert copied.single.owner == 100 + assert copied.tracer.trees[0].parameters[-1].name == "rate.single.owner" + assert node.tracer.trees[0].parameters == [] + + +@pytest.mark.parametrize("copier", COPIERS.values(), ids=COPIERS.keys()) +def test_copied_tracing_wrapper_of_a_vectorial_node_keeps_both_types(copier): + vectorial = rate_node()("2015-01-01")[np.array(["couple", "single"])] + node = TracingParameterNodeAtInstant(vectorial, FullTracer(), "default") + + copied = copier(node) + + assert type(copied) is TracingParameterNodeAtInstant + assert type(copied.parameter_node_at_instant) is VectorialParameterNodeAtInstant + assert_near(copied.parameter_node_at_instant.owner, [500, 100]) + + +@pytest.mark.parametrize( + "wrapper_class", [VectorialParameterNodeAtInstant, TracingParameterNodeAtInstant] +) +def test_lookups_on_an_unfilled_parameter_wrapper_raise_attribute_error( + wrapper_class, +): + unfilled = wrapper_class.__new__(wrapper_class) + + assert getattr(unfilled, "__setstate__", None) is None + assert not hasattr(unfilled, "anything") diff --git a/tests/core/test_simulation_copy_pickle_property.py b/tests/core/test_simulation_copy_pickle_property.py new file mode 100644 index 00000000..84dea200 --- /dev/null +++ b/tests/core/test_simulation_copy_pickle_property.py @@ -0,0 +1,99 @@ +"""A copied simulation is a simulation in its own right. + +For random situations, warm caches and branches, a ``copy.deepcopy`` or +pickle round trip of a simulation calculates what a freshly built simulation +does, and writing to the copy leaves the original unchanged. +``test_simulation_copy_pickle.py`` pins the same behaviour with examples. +""" + +from __future__ import annotations + +import pytest + +# The country-package smoke job installs no dev dependencies. +hypothesis = pytest.importorskip("hypothesis") +st = hypothesis.strategies + +from policyengine_core.country_template import CountryTaxBenefitSystem +from policyengine_core.tools import assert_near +from tests.fixtures.simulation_copy import COPIERS, FEB, JAN, build_simulation + +YEAR = "2025" +VARIABLES = ( + "income_tax", + "social_security_contribution", + "basic_income", + "disposable_income", + "household_income", + "total_taxes", + "total_benefits", +) + + +@st.composite +def situations(draw): + salaries = draw(st.lists(st.integers(0, 20_000), min_size=1, max_size=4)) + count = len(salaries) + order = draw(st.permutations(range(count))) + cuts = sorted(draw(st.sets(st.integers(1, count - 1)))) if count > 1 else [] + bounds = [0, *cuts, count] + households = [list(order[a:b]) for a, b in zip(bounds, bounds[1:])] + occupancy = draw( + st.lists( + st.sampled_from([None, "owner", "tenant", "free_lodger", "homeless"]), + min_size=len(households), + max_size=len(households), + ) + ) + warmed = draw(st.lists(st.sampled_from(VARIABLES), unique=True)) + branch_salary = draw(st.none() | st.integers(0, 20_000)) + return salaries, households, occupancy, warmed, branch_salary + + +@hypothesis.settings( + max_examples=60, + deadline=None, + suppress_health_check=[hypothesis.HealthCheck.too_slow], +) +@hypothesis.given(situation=situations(), copier_name=st.sampled_from(sorted(COPIERS))) +def test_copy_calculates_like_a_fresh_simulation_and_stays_independent( + situation, copier_name +): + """For any situation and warm cache, a copy matches a freshly built + simulation on every variable, and writing to the copy leaves the original + unchanged.""" + salaries, households, occupancy, warmed, branch_salary = situation + tax_benefit_system = CountryTaxBenefitSystem() + original = build_simulation(tax_benefit_system, salaries, households, occupancy) + for variable in warmed: + original.calculate(variable, JAN) + if branch_salary is not None: + original.get_branch("child").set_input( + "salary", FEB, [branch_salary] * len(salaries) + ) + + copied = COPIERS[copier_name](original) + fresh = build_simulation(tax_benefit_system, salaries, households, occupancy) + + for variable in VARIABLES: + assert_near(copied.calculate(variable, JAN), fresh.calculate(variable, JAN)) + assert_near( + copied.calculate("housing_tax", YEAR), fresh.calculate("housing_tax", YEAR) + ) + assert list(copied.household("housing_occupancy_status", JAN).decode_to_str()) == [ + status or "tenant" for status in occupancy + ] + if branch_salary is not None: + fresh_child = fresh.get_branch("child") + fresh_child.set_input("salary", FEB, [branch_salary] * len(salaries)) + assert_near( + copied.branches["child"].calculate("income_tax", FEB), + fresh_child.calculate("income_tax", FEB), + ) + + copied.set_input("salary", FEB, [1] * len(salaries)) + copied.calculate("disposable_income", FEB) + assert original.persons.get_holder("salary").get_array(FEB) is None + assert original.persons.get_holder("disposable_income").get_array(FEB) is None + for variable in warmed: + assert_near(original.calculate(variable, JAN), fresh.calculate(variable, JAN)) diff --git a/tests/fixtures/simulation_copy.py b/tests/fixtures/simulation_copy.py new file mode 100644 index 00000000..596e4672 --- /dev/null +++ b/tests/fixtures/simulation_copy.py @@ -0,0 +1,81 @@ +"""Helpers for the simulation copy and pickle tests.""" + +from __future__ import annotations + +import copy +import pickle + +from policyengine_core.country_template import CountryTaxBenefitSystem +from policyengine_core.parameters import ParameterNode +from policyengine_core.reforms import Reform +from policyengine_core.simulations import SimulationBuilder + +JAN = "2025-01" +FEB = "2025-02" + +COPIERS = { + "deepcopy": copy.deepcopy, + "pickle": lambda value: pickle.loads(pickle.dumps(value)), +} + + +def build_simulation( + tax_benefit_system=None, salaries=(1000,), households=None, occupancy=None +): + """One household per entry of ``households`` (lists of person indices). + + ``occupancy`` gives each household's ``housing_occupancy_status``. + """ + tax_benefit_system = tax_benefit_system or CountryTaxBenefitSystem() + names = [f"p{i}" for i in range(len(salaries))] + households = households or [list(range(len(salaries)))] + occupancy = occupancy or [None] * len(households) + return SimulationBuilder().build_from_entities( + tax_benefit_system, + { + "persons": { + name: {"salary": {JAN: salary}} for name, salary in zip(names, salaries) + }, + "households": { + f"h{i}": { + "parents": [names[j] for j in members[:2]], + "children": [names[j] for j in members[2:]], + "accommodation_size": {JAN: 100}, + **( + {"housing_occupancy_status": {JAN: occupancy[i]}} + if occupancy[i] + else {} + ), + } + for i, members in enumerate(households) + }, + }, + ) + + +def rate_node() -> ParameterNode: + return ParameterNode( + "rate", + data={ + "single": { + "owner": {"values": {"2015-01-01": 100}}, + "tenant": {"values": {"2015-01-01": 300}}, + }, + "couple": { + "owner": {"values": {"2015-01-01": 500}}, + "tenant": {"values": {"2015-01-01": 700}}, + }, + }, + ) + + +class DoubleIncomeTaxRate(Reform): + """Module level, so pickle can find the reform class by name.""" + + def apply(self): + def double(parameters): + rate = parameters.taxes.income_tax_rate + rate.update(period="year:2015:20", value=rate("2015-01-01") * 2) + return parameters + + self.modify_parameters(double)