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 000000000..e4c0e9bf3 --- /dev/null +++ b/changelog.d/fix-simulation-copy-pickle-recursion.fixed.md @@ -0,0 +1 @@ +Fix simulation copying and same-process unpickling, preserve `EnumArray` subclasses and enum metadata through pickling, and allow legacy enum arrays to be re-pickled. diff --git a/policyengine_core/enums/enum_array.py b/policyengine_core/enums/enum_array.py index c518add8c..789f48613 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,37 @@ from policyengine_core.enums import Enum +def _restore_enum_array( + array: numpy.ndarray, + enum_name: Optional[Tuple[str, str]], + array_type: Optional[Type[EnumArray]] = None, +) -> 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 + # Do not invoke a subclass's constructor: it may require other arguments. + # The default also accepts pickles written by the two-argument reducer. + restored = array.view(EnumArray if array_type is None else array_type) + restored.possible_values = possible_values + return restored + + class EnumArray(numpy.ndarray): """ Numpy array subclass representing an array of enum items. @@ -35,6 +68,16 @@ 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``. + # Legacy NumPy pickles do not restore this attribute at all. + enum = getattr(self, "possible_values", None) + name = None if enum is None else (enum.__module__, enum.__qualname__) + return _restore_enum_array, (self.view(numpy.ndarray), name, type(self)) + 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/populations/population.py b/policyengine_core/populations/population.py index d2e8df1da..9d64d1ef3 100644 --- a/policyengine_core/populations/population.py +++ b/policyengine_core/populations/population.py @@ -53,6 +53,28 @@ 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. + # Reject protocol probes before shortcut lookup, including when a + # subclass's entity property reads state that has not been restored. + if attribute in ( + "entity", + "simulation", + "__slots__", + "__copy__", + "__deepcopy__", + "__getstate__", + "__setstate__", + "__reduce__", + "__reduce_ex__", + "__getnewargs__", + "__getnewargs_ex__", + ): + 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/tests/core/test_simulation_copy_pickle.py b/tests/core/test_simulation_copy_pickle.py new file mode 100644 index 000000000..c0a644b59 --- /dev/null +++ b/tests/core/test_simulation_copy_pickle.py @@ -0,0 +1,454 @@ +"""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 ( + ParameterNode, + ParameterNodeAtInstant, + 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, + ChildEnumArray, + DoubleIncomeTaxRate, + build_simulation, + rate_node, +) + + +class PropertyBackedPopulation(Population): + @property + def entity(self): + return self._entity + + @entity.setter + def entity(self, value): + self._entity = value + + +class SlottedNodeState: + __slots__ = () + + def __getstate__(self): + return self.__dict__, self.marker + + def __setstate__(self, state): + attributes, self.marker = state + self.__dict__.update(attributes) + + +class SlottedScalarNode(SlottedNodeState, ParameterNodeAtInstant): + __slots__ = ("marker",) + + +class SlottedVectorialNode(SlottedNodeState, VectorialParameterNodeAtInstant): + __slots__ = ("marker",) + + +# 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 + + +@pytest.mark.parametrize("copier", [copy.copy, copy.deepcopy]) +def test_property_backed_population_can_be_copied(copier): + population = PropertyBackedPopulation(build_simulation().persons.entity) + population.ids = ["person"] + + copied = copier(population) + + assert type(copied) is PropertyBackedPopulation + assert copied.entity.key == population.entity.key + assert copied.ids == population.ids + + +@pytest.mark.parametrize("protocol", range(6)) +def test_property_backed_population_can_be_pickled(protocol): + population = PropertyBackedPopulation(build_simulation().persons.entity) + population.ids = ["person"] + + restored = pickle.loads(pickle.dumps(population, protocol=protocol)) + + assert type(restored) is PropertyBackedPopulation + assert restored.entity.key == population.entity.key + assert restored.ids == population.ids + + +# 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)) +@pytest.mark.parametrize("array_class", [EnumArray, ChildEnumArray]) +def test_pickled_enum_array_keeps_its_possible_values(protocol, array_class): + status = build_simulation(occupancy=["free_lodger"]).household( + "housing_occupancy_status", JAN + ) + status = status.view(array_class) + + restored = pickle.loads(pickle.dumps(status, protocol=protocol)) + + assert type(restored) is array_class + 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] + + +@pytest.mark.parametrize("protocol", range(6)) +@pytest.mark.parametrize("array_class", [EnumArray, ChildEnumArray]) +def test_enum_array_pickled_by_an_earlier_release_still_loads(protocol, array_class): + """Earlier releases pickled an EnumArray as numpy pickles any ndarray.""" + status = build_simulation(occupancy=["owner"]).household( + "housing_occupancy_status", JAN + ) + status = status.view(array_class) + + 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, protocol=protocol).dump(status) + + restored = pickle.loads(buffer.getvalue()) + + assert type(restored) is array_class + assert restored.tolist() == status.tolist() + assert not hasattr(restored, "possible_values") + + repickled = pickle.loads(pickle.dumps(restored, protocol=protocol)) + + assert type(repickled) is array_class + assert repickled.possible_values is None + assert repickled.dtype == status.dtype + assert repickled.shape == status.shape + np.testing.assert_array_equal(repickled.view(np.ndarray), status.view(np.ndarray)) + + +@pytest.mark.parametrize("array_class", [EnumArray, ChildEnumArray]) +def test_copied_enum_array_keeps_its_possible_values(array_class): + status = build_simulation(occupancy=["owner"]).household( + "housing_occupancy_status", JAN + ) + status = status.view(array_class) + + for copied in (copy.copy(status), copy.deepcopy(status)): + assert type(copied) is array_class + 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. + + +def test_vectorial_parameter_field_named_vector_still_resolves(): + rate = ParameterNode( + "rate", + data={ + "single": {"vector": {"values": {"2015-01-01": 100}}}, + "couple": {"vector": {"values": {"2015-01-01": 500}}}, + }, + ) + vectorial = rate("2015-01-01")[np.array(["single", "couple"])] + + np.testing.assert_array_equal(vectorial["vector"], [100, 500]) + + +@pytest.mark.parametrize("protocol", range(6)) +@pytest.mark.parametrize("node_class", [SlottedScalarNode, SlottedVectorialNode]) +def test_tracing_wrapper_of_a_slotted_node_can_be_pickled(protocol, node_class): + source = rate_node()("2015-01-01") + if node_class is SlottedVectorialNode: + source = source[np.array(["single", "couple"])] + inner = node_class.__new__(node_class) + inner.__dict__.update(source.__dict__) + inner.marker = "slotted" + + # The wrapper must preserve the protocol support of its wrapped node. + unwrapped = pickle.loads(pickle.dumps(inner, protocol=protocol)) + wrapped = TracingParameterNodeAtInstant(inner, FullTracer(), "default") + restored = pickle.loads(pickle.dumps(wrapped, protocol=protocol)) + + assert type(restored) is TracingParameterNodeAtInstant + restored_inner = restored.parameter_node_at_instant + assert type(restored_inner) is type(unwrapped) is node_class + assert restored_inner.marker == unwrapped.marker == "slotted" + if node_class is SlottedVectorialNode: + np.testing.assert_array_equal(restored_inner.owner, unwrapped.owner) + else: + assert restored_inner.single.owner == unwrapped.single.owner == 100 + + +@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 000000000..5a3d02eda --- /dev/null +++ b/tests/core/test_simulation_copy_pickle_property.py @@ -0,0 +1,138 @@ +"""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 copy +import pickle + +import numpy as np +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.enums import EnumArray +from policyengine_core.tools import assert_near +from tests.fixtures.simulation_copy import ( + COPIERS, + FEB, + JAN, + ChildEnumArray, + CopyEnum, + 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)) + + +@hypothesis.settings(max_examples=40, deadline=None) +@hypothesis.given( + indices=st.lists(st.integers(0, 1), max_size=40), + dtype=st.sampled_from(["uint8", "int16", "int32"]), + array_class=st.sampled_from([EnumArray, ChildEnumArray]), + protocol=st.integers(0, 5), +) +def test_enum_array_copy_and_pickle_preserve_numpy_values( + indices, dtype, array_class, protocol +): + """Copies and round trips agree with the plain NumPy array reference.""" + reference = np.array(indices, dtype=dtype).reshape(-1, 1) + original = array_class(reference.copy(), CopyEnum) + + for restored in ( + copy.copy(original), + copy.deepcopy(original), + pickle.loads(pickle.dumps(original, protocol=protocol)), + ): + assert type(restored) is array_class + assert restored.possible_values is CopyEnum + assert restored.dtype == reference.dtype + assert restored.shape == reference.shape + np.testing.assert_array_equal(restored.view(np.ndarray), reference) + assert not np.shares_memory(restored, original) diff --git a/tests/fixtures/simulation_copy.py b/tests/fixtures/simulation_copy.py new file mode 100644 index 000000000..d89dde091 --- /dev/null +++ b/tests/fixtures/simulation_copy.py @@ -0,0 +1,91 @@ +"""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.enums import Enum, EnumArray +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)), +} + + +class CopyEnum(Enum): + first = "First" + second = "Second" + + +class ChildEnumArray(EnumArray): + """Importable subclass for copy and pickle type-preservation checks.""" + + +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)