From b6c68eb363fd746b60e2a7c8dacfefcbd622840e Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Fri, 2 Oct 2026 09:17:13 -0400 Subject: [PATCH 1/2] 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 000000000..1332e922e --- /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 c518add8c..36be3f03e 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 fcc562ef0..81f1a9d3c 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 d2e8df1da..d391b8105 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 d4bf17466..5d3fde49b 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 000000000..89fc212ad --- /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 000000000..84dea2002 --- /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 000000000..596e4672d --- /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) From f310a8e79d3c27f8c56d41443450c26edb27341a Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Wed, 7 Oct 2026 17:01:00 -0400 Subject: [PATCH 2/2] Preserve EnumArray subclasses and legacy re-pickling (#567) Guard Population protocol probes before projector lookup. Add regressions for protocols 0-5, property-backed populations, slotted tracing nodes, vector fields, and NumPy differential copy behavior. Co-Authored-By: Claude Opus 5.5 --- ...-simulation-copy-pickle-recursion.fixed.md | 2 +- policyengine_core/enums/enum_array.py | 15 +- policyengine_core/populations/population.py | 16 ++- tests/core/test_simulation_copy_pickle.py | 129 ++++++++++++++++-- .../test_simulation_copy_pickle_property.py | 41 +++++- tests/fixtures/simulation_copy.py | 10 ++ 6 files changed, 198 insertions(+), 15 deletions(-) diff --git a/changelog.d/fix-simulation-copy-pickle-recursion.fixed.md b/changelog.d/fix-simulation-copy-pickle-recursion.fixed.md index 1332e922e..e4c0e9bf3 100644 --- a/changelog.d/fix-simulation-copy-pickle-recursion.fixed.md +++ b/changelog.d/fix-simulation-copy-pickle-recursion.fixed.md @@ -1 +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. +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 36be3f03e..789f48613 100644 --- a/policyengine_core/enums/enum_array.py +++ b/policyengine_core/enums/enum_array.py @@ -12,7 +12,9 @@ def _restore_enum_array( - array: numpy.ndarray, enum_name: Optional[Tuple[str, str]] + 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. @@ -33,7 +35,11 @@ def _restore_enum_array( possible_values = operator.attrgetter(qualified_name)(module) except (ImportError, AttributeError): pass - return EnumArray(array, possible_values) + # 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): @@ -67,9 +73,10 @@ def __reduce__(self) -> tuple: # ``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 + # 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) + 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 diff --git a/policyengine_core/populations/population.py b/policyengine_core/populations/population.py index d391b8105..9d64d1ef3 100644 --- a/policyengine_core/populations/population.py +++ b/policyengine_core/populations/population.py @@ -57,7 +57,21 @@ def __getattr__(self, attribute: str) -> Any: # 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"): + # 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}" ) diff --git a/tests/core/test_simulation_copy_pickle.py b/tests/core/test_simulation_copy_pickle.py index 89fc212ad..c0a644b59 100644 --- a/tests/core/test_simulation_copy_pickle.py +++ b/tests/core/test_simulation_copy_pickle.py @@ -28,7 +28,11 @@ from policyengine_core.country_template import CountryTaxBenefitSystem from policyengine_core.enums import Enum, EnumArray -from policyengine_core.parameters import VectorialParameterNodeAtInstant +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 @@ -36,11 +40,42 @@ 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. @@ -82,6 +117,30 @@ def test_shallow_copy_of_a_population_shares_its_holders(): 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. @@ -195,14 +254,16 @@ def test_copied_simulation_keeps_enum_inputs(copier): @pytest.mark.parametrize("protocol", range(pickle.HIGHEST_PROTOCOL + 1)) -def test_pickled_enum_array_keeps_its_possible_values(protocol): +@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 EnumArray + 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() @@ -244,11 +305,14 @@ class Unreachable(Enum): assert restored.tolist() == [1, 0, 1] -def test_enum_array_pickled_by_an_earlier_release_still_loads(): +@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): @@ -257,21 +321,32 @@ def reducer_override(self, obj): return NotImplemented buffer = io.BytesIO() - EarlierPickler(buffer).dump(status) + EarlierPickler(buffer, protocol=protocol).dump(status) restored = pickle.loads(buffer.getvalue()) - assert type(restored) is EnumArray + 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)) -def test_copied_enum_array_keeps_its_possible_values(): +@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 EnumArray + 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) @@ -280,6 +355,44 @@ def test_copied_enum_array_keeps_its_possible_values(): # 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"])] diff --git a/tests/core/test_simulation_copy_pickle_property.py b/tests/core/test_simulation_copy_pickle_property.py index 84dea2002..5a3d02eda 100644 --- a/tests/core/test_simulation_copy_pickle_property.py +++ b/tests/core/test_simulation_copy_pickle_property.py @@ -8,6 +8,10 @@ from __future__ import annotations +import copy +import pickle + +import numpy as np import pytest # The country-package smoke job installs no dev dependencies. @@ -15,8 +19,16 @@ 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, build_simulation +from tests.fixtures.simulation_copy import ( + COPIERS, + FEB, + JAN, + ChildEnumArray, + CopyEnum, + build_simulation, +) YEAR = "2025" VARIABLES = ( @@ -97,3 +109,30 @@ def test_copy_calculates_like_a_fresh_simulation_and_stays_independent( 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 index 596e4672d..d89dde091 100644 --- a/tests/fixtures/simulation_copy.py +++ b/tests/fixtures/simulation_copy.py @@ -6,6 +6,7 @@ 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 @@ -19,6 +20,15 @@ } +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 ):