From 98d37ea0bee84ad4743b9a12f7f6e7a0dc75f316 Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Sat, 3 Oct 2026 10:14:32 -0400 Subject: [PATCH] Let vectorial parameter nodes be copied and pickled VectorialParameterNodeAtInstant.__getattr__ delegated every missing attribute to its recarray. copy and pickle probe a new, empty instance for __setstate__, so the lookup of self.vector recursed; deepcopy found the vector's __deepcopy__ and returned a bare recarray. A reform that calls modify_parameters deep-copies the baseline tree, cached fancy-indexing nodes included, and got a recarray in their place. Mirror #574's guard: answer the copy and pickle protocol names (and __slots__) and any lookup on an unfilled instance with AttributeError, still delegating NumPy's array protocol. Build a child read by name with the node's name and instant (it passed one of three required arguments), and leave the id-keyed enum lookup cache out of copied state. Co-Authored-By: Claude Opus 5.5 --- .../fix-vectorial-node-copy-pickle.fixed.md | 1 + .../vectorial_parameter_node_at_instant.py | 57 ++++- .../test_vectorial_parameter_node_copy.py | 219 ++++++++++++++++++ ..._vectorial_parameter_node_copy_property.py | 132 +++++++++++ tests/fixtures/vectorial_parameter_nodes.py | 76 ++++++ 5 files changed, 483 insertions(+), 2 deletions(-) create mode 100644 changelog.d/fix-vectorial-node-copy-pickle.fixed.md create mode 100644 tests/core/parameters/test_vectorial_parameter_node_copy.py create mode 100644 tests/core/parameters/test_vectorial_parameter_node_copy_property.py create mode 100644 tests/fixtures/vectorial_parameter_nodes.py diff --git a/changelog.d/fix-vectorial-node-copy-pickle.fixed.md b/changelog.d/fix-vectorial-node-copy-pickle.fixed.md new file mode 100644 index 000000000..8d4615cc1 --- /dev/null +++ b/changelog.d/fix-vectorial-node-copy-pickle.fixed.md @@ -0,0 +1 @@ +Let vectorial parameter nodes be copied, deep-copied and pickled, so a reform that calls `modify_parameters` keeps a cached fancy-indexing node as a node, and stop reading a child node by name (`node.owner`) from raising `TypeError`. diff --git a/policyengine_core/parameters/vectorial_parameter_node_at_instant.py b/policyengine_core/parameters/vectorial_parameter_node_at_instant.py index fcc562ef0..ea7df3d2a 100644 --- a/policyengine_core/parameters/vectorial_parameter_node_at_instant.py +++ b/policyengine_core/parameters/vectorial_parameter_node_at_instant.py @@ -51,6 +51,47 @@ def _unify_structured_dtypes(values): return unified_dtype, all_fields, casted_values +# Copying and pickling look these names up on an instance. Answered by the +# vector, they would copy or restore the vector, not the node. Pickle +# protocols 0 and 1 read ``__slots__`` to decide whether the node can be +# pickled at all. +_COPY_PROTOCOL = frozenset( + { + "__copy__", + "__deepcopy__", + "__getstate__", + "__setstate__", + "__reduce__", + "__reduce_ex__", + "__getnewargs__", + "__getnewargs_ex__", + "__slots__", + } +) + + +def _wrapped(wrapper: object, attribute: str, key: str) -> object: + """Return the object ``wrapper`` delegates the lookup of ``key`` to. + + ``__getattr__`` only runs when normal lookup fails. Two such lookups must + not reach the wrapped object: + + - the copy and pickle protocol (``__deepcopy__``, ``__setstate__``, ...); + - any name on an instance ``copy`` or ``pickle`` has created with + ``__new__`` and not filled in yet. It has no ``attribute``, so reading + it here would call ``__getattr__`` again, without end. + + Every other name is delegated, special names included: NumPy reads its + array protocol (``__array_struct__``, ...) from a vectorial node this way. + """ + if key in _COPY_PROTOCOL: + raise AttributeError(key) + try: + return wrapper.__dict__[attribute] + except KeyError: + raise AttributeError(key) from None + + class VectorialParameterNodeAtInstant: """ Parameter node of the legislation at a given instant which has been vectorized. @@ -209,11 +250,23 @@ def __init__(self, name: str, vector: ArrayLike, instant_str: str): self._instant_str = instant_str def __getattr__(self, attribute: str) -> Any: - result = getattr(self.vector, attribute) + result = getattr(_wrapped(self, "vector", attribute), attribute) if isinstance(result, numpy.recarray): - return VectorialParameterNodeAtInstant(result) + # A child that is itself a node, such as ``node.owner`` when + # ``owner`` has children. Like ``__getitem__``, keep this node's + # name and instant. + return VectorialParameterNodeAtInstant( + self._name, result, self._instant_str + ) return result + def __getstate__(self) -> dict: + # ``_enum_lut_cache`` is keyed by ``id(enum)``. In another process that + # id can belong to a different enum, so copies rebuild the cache. + state = self.__dict__.copy() + state.pop("_enum_lut_cache", None) + return state + def __getitem__(self, key: str) -> Any: # If the key is a string, just get the subnode if isinstance(key, str): diff --git a/tests/core/parameters/test_vectorial_parameter_node_copy.py b/tests/core/parameters/test_vectorial_parameter_node_copy.py new file mode 100644 index 000000000..98af4f9b2 --- /dev/null +++ b/tests/core/parameters/test_vectorial_parameter_node_copy.py @@ -0,0 +1,219 @@ +"""Vectorial parameter nodes survive ``copy``, ``copy.deepcopy`` and pickle. + +``VectorialParameterNodeAtInstant.__getattr__`` hands every attribute the node +lacks to its ``vector``, a ``numpy.recarray``. ``copy`` and ``pickle`` create +an empty instance and probe it for ``__setstate__`` before they restore its +``__dict__``. With no ``vector`` yet, ``__getattr__`` looked ``vector`` up +through itself until ``RecursionError``. ``copy.deepcopy`` looks +``__deepcopy__`` up on the instance, got the vector's, and returned a bare +``recarray``. + +``ParameterNodeAtInstant`` caches its vectorial node, and +``Reform.modify_parameters`` deep-copies the baseline's parameter tree, so a +reform found a ``recarray`` where its cached node should be. +""" + +import copy + +import numpy as np +import pytest + +from policyengine_core.enums import EnumArray +from policyengine_core.parameters import ( + ParameterNotFoundError, + VectorialParameterNodeAtInstant, +) +from policyengine_core.reforms import Reform +from policyengine_core.tools import assert_near +from tests.fixtures.vectorial_parameter_nodes import ( + COPIERS, + DEEP_COPIERS, + INSTANT, + Family, + SlottedRecords, + fancy_indexing_parameters, +) + + +def single_rate_node() -> VectorialParameterNodeAtInstant: + return fancy_indexing_parameters()(INSTANT).rate.single[ + np.array(["owner", "tenant"]) + ] + + +@pytest.mark.parametrize("copier", COPIERS.values(), ids=COPIERS.keys()) +def test_copy_is_a_vectorial_node_with_the_same_contents(copier): + node = single_rate_node() + + copied = copier(node) + + assert type(copied) is VectorialParameterNodeAtInstant + assert copied._name == node._name == "rate.single" + assert copied._instant_str == node._instant_str == INSTANT + assert copied.vector.dtype == node.vector.dtype + assert copied.vector.tobytes() == node.vector.tobytes() + assert_near(copied.z1, [100, 300]) + assert_near(copied[np.array(["z2", "z1"])], [200, 300]) + + +@pytest.mark.parametrize("copier", DEEP_COPIERS.values(), ids=DEEP_COPIERS.keys()) +def test_deep_copy_has_a_vector_of_its_own(copier): + node = single_rate_node() + + copied = copier(node) + copied.vector.z1[:] = 0 + + assert not np.shares_memory(copied.vector, node.vector) + assert_near(node.z1, [100, 300]) + + +def test_shallow_copy_shares_the_vector(): + node = single_rate_node() + + assert copy.copy(node).vector is node.vector + + +# ``object`` has no such attribute, so a lookup reaches ``__getattr__``. +@pytest.mark.parametrize( + "name", + [ + "__copy__", + "__deepcopy__", + "__setstate__", + "__getnewargs__", + "__getnewargs_ex__", + "__slots__", + ], +) +def test_copy_protocol_is_not_read_from_the_vector(name): + assert not hasattr(single_rate_node(), name) + + +# ``object`` has these, so a lookup never reaches ``__getattr__``. +@pytest.mark.parametrize("name", ["__getstate__", "__reduce__", "__reduce_ex__"]) +def test_copy_protocol_methods_are_the_nodes_own(name): + node = single_rate_node() + + assert getattr(node, name).__self__ is node + + +@pytest.mark.parametrize("copier", COPIERS.values(), ids=COPIERS.keys()) +def test_node_with_a_slotted_vector_is_copied(copier): + """The node's own ``__getstate__`` keeps the vector's ``__slots__`` out.""" + vector = single_rate_node().vector.view(SlottedRecords) + node = VectorialParameterNodeAtInstant("rate.single", vector, INSTANT) + assert vector.__slots__ == ("note",) + + copied = copier(node) + + assert type(copied) is VectorialParameterNodeAtInstant + assert copied.vector.tobytes() == vector.tobytes() + assert_near(copied.z1, [100, 300]) + + +@pytest.mark.parametrize( + "name", ["vector", "_name", "dtype", "__setstate__", "__array_struct__"] +) +def test_lookups_on_an_unfilled_node_raise_attribute_error(name): + """``copy`` and ``pickle`` probe a node before they fill it in.""" + unfilled = VectorialParameterNodeAtInstant.__new__(VectorialParameterNodeAtInstant) + + with pytest.raises(AttributeError): + getattr(unfilled, name) + + +def test_numpy_reads_the_vector_through_the_node(): + node = single_rate_node() + + array = np.asarray(node) + + assert type(array) is np.ndarray + assert array.shape == node.vector.shape + assert array.dtype.itemsize == node.vector.dtype.itemsize + assert array.tobytes() == node.vector.tobytes() + assert node.dtype == node.vector.dtype + assert node.shape == node.vector.shape + + +@pytest.mark.parametrize("copier", COPIERS.values(), ids=COPIERS.keys()) +def test_numpy_reads_a_copy_as_it_reads_the_original(copier): + node = single_rate_node() + + original, copied = np.asarray(node), np.asarray(copier(node)) + + assert copied.dtype == original.dtype + assert copied.shape == original.shape + assert copied.tobytes() == original.tobytes() + + +def test_child_node_read_by_name_is_a_vectorial_node(): + """``__getattr__`` built the child with one of its three arguments.""" + node = fancy_indexing_parameters()(INSTANT).rate[ + np.array(["single", "couple", "single"]) + ] + + for owner in (node.owner, node["owner"]): + assert type(owner) is VectorialParameterNodeAtInstant + assert owner._name == node._name == "rate" + assert owner._instant_str == INSTANT + assert_near(owner.z1, [100, 500, 100]) + assert_near(owner[np.array(["z2", "z1", "z1"])], [200, 500, 100]) + with pytest.raises(ParameterNotFoundError): + owner[np.array(["z1", "z9", "z1"])] + + +@pytest.mark.parametrize("copier", COPIERS.values(), ids=COPIERS.keys()) +def test_copy_reads_enum_keys_as_the_original_does(copier): + node = VectorialParameterNodeAtInstant.build_from_node( + fancy_indexing_parameters()(INSTANT).rate + ) + families = EnumArray(np.array([1, 0, 1]), Family) + owners = np.array(["owner", "owner", "owner"]) + expected = node[families][owners].z2 # Fills the node's enum lookup cache. + + copied = copier(node) + + # The cache is keyed by ``id(enum)``, which in another process can belong + # to a different enum. + assert "_enum_lut_cache" not in copied.__dict__ + assert_near(expected, [600, 200, 600]) + assert_near(copied[families][owners].z2, expected) + family_items = np.array([Family.couple, Family.single, Family.couple], dtype=object) + assert_near(copied[family_items][owners].z2, expected) + + +@pytest.mark.parametrize("copier", DEEP_COPIERS.values(), ids=DEEP_COPIERS.keys()) +def test_copied_parameter_tree_keeps_its_cached_vectorial_node(copier): + parameters = fancy_indexing_parameters() + families = np.array(["single", "couple", "couple"]) + housings = np.array(["owner", "tenant", "owner"]) + expected = parameters(INSTANT).rate[families][housings].z2 + assert type(parameters(INSTANT).rate._vectorial_node) is ( + VectorialParameterNodeAtInstant + ) + + copied = copier(parameters) + + rate = copied(INSTANT).rate + assert type(rate._vectorial_node) is VectorialParameterNodeAtInstant + assert_near(expected, [200, 800, 600]) + assert_near(rate[families][housings].z2, expected) + + +def test_reform_that_modifies_parameters_keeps_the_cached_vectorial_node( + isolated_tax_benefit_system, +): + system = isolated_tax_benefit_system + names = np.array(["age_of_retirement", "age_of_majority", "age_of_retirement"]) + # Caches the vectorial node of ``general`` at this instant. + expected = system.get_parameters_at_instant("2017-01-01").general[names] + + class KeepParameters(Reform): + def apply(self): + # Deep-copies the baseline's tree, cache included. + self.modify_parameters(lambda parameters: parameters) + + general = KeepParameters(system).get_parameters_at_instant("2017-01-01").general + + assert type(general._vectorial_node) is VectorialParameterNodeAtInstant + assert_near(general[names], expected) diff --git a/tests/core/parameters/test_vectorial_parameter_node_copy_property.py b/tests/core/parameters/test_vectorial_parameter_node_copy_property.py new file mode 100644 index 000000000..100c04403 --- /dev/null +++ b/tests/core/parameters/test_vectorial_parameter_node_copy_property.py @@ -0,0 +1,132 @@ +"""A copy of a vectorial parameter node reads what the original reads. + +Random paths through ``rate...`` take each level by +name or by an array of keys (strings, an ``EnumArray`` or enum items). Partway +along, the vectorial node reached so far is copied with ``copy``, +``copy.deepcopy`` or a pickle round trip, and the rest of the path is read +from the copy and from the original. ``RATE`` gives the values each read must +return. ``test_vectorial_parameter_node_copy.py`` pins the same behaviour with +examples. +""" + +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.enums import EnumArray +from policyengine_core.parameters import VectorialParameterNodeAtInstant +from policyengine_core.tools import assert_near +from tests.fixtures.vectorial_parameter_nodes import ( + COPIERS, + DEEP_COPIERS, + FAMILIES, + HOUSINGS, + INSTANT, + RATE, + ZONES, + Family, + Housing, + fancy_indexing_parameters, +) + +LEVELS = (FAMILIES, HOUSINGS, ZONES) +ENUMS = (Family, Housing, None) + + +@st.composite +def _paths(draw): + """Steps for the three levels, and how many to take before copying.""" + size = draw(st.integers(1, 6)) + steps = [] + for children, enum in zip(LEVELS, ENUMS): + if draw(st.booleans()): + # One child for every element, read as an attribute or an item. + steps.append( + ( + "name", + draw(st.sampled_from(children)), + draw(st.sampled_from(["attribute", "item"])), + ) + ) + else: + keys = draw( + st.lists(st.sampled_from(children), min_size=size, max_size=size) + ) + kinds = ["strings"] + (["enum_array", "enum_items"] if enum else []) + steps.append(("keys", keys, draw(st.sampled_from(kinds)))) + # Copy a vectorial node: after an array of keys, before the leaf. + first_keys = next( + (level for level, step in enumerate(steps) if step[0] == "keys"), None + ) + hypothesis.assume(first_keys is not None and first_keys < 2) + copy_after = draw(st.integers(first_keys + 1, 2)) + return size, steps, copy_after + + +def _read(node, step, enum): + if step[0] == "name": + _, child, how = step + return getattr(node, child) if how == "attribute" else node[child] + _, keys, kind = step + if kind == "strings": + return node[np.array(keys)] + items = [enum[key] for key in keys] + if kind == "enum_items": + return node[np.array(items, dtype=object)] + return node[EnumArray(np.array([item.index for item in items]), enum)] + + +def _read_path(node, steps, levels): + for level in levels: + node = _read(node, steps[level], ENUMS[level]) + return node + + +def _expected(size, steps): + def key(level, element): + step = steps[level] + return step[1] if step[0] == "name" else step[1][element] + + return [ + RATE[key(0, element)][key(1, element)][key(2, element)] + for element in range(size) + ] + + +@hypothesis.settings(max_examples=300, deadline=None) +@hypothesis.given( + path=_paths(), + copier_name=st.sampled_from(sorted(COPIERS)), + read_before_copying=st.booleans(), +) +def test_copy_reads_what_the_original_reads(path, copier_name, read_before_copying): + size, steps, copy_after = path + node = _read_path( + fancy_indexing_parameters()(INSTANT).rate, steps, range(copy_after) + ) + assert type(node) is VectorialParameterNodeAtInstant + rest = range(copy_after, len(LEVELS)) + expected = _expected(size, steps) + if read_before_copying: + # Fills the node's lookup caches before it is copied. + assert_near(_read_path(node, steps, rest), expected) + + copied = COPIERS[copier_name](node) + + assert type(copied) is VectorialParameterNodeAtInstant + assert copied._name == node._name + assert copied._instant_str == node._instant_str + assert copied.vector.dtype == node.vector.dtype + assert copied.vector.tobytes() == node.vector.tobytes() + original_array, copied_array = np.asarray(node), np.asarray(copied) + assert copied_array.dtype == original_array.dtype + assert copied_array.tobytes() == original_array.tobytes() + if copier_name in DEEP_COPIERS: + assert not np.shares_memory(copied.vector, node.vector) + else: + assert copied.vector is node.vector + assert_near(_read_path(copied, steps, rest), expected) + assert_near(_read_path(node, steps, rest), expected) diff --git a/tests/fixtures/vectorial_parameter_nodes.py b/tests/fixtures/vectorial_parameter_nodes.py new file mode 100644 index 000000000..5e971728d --- /dev/null +++ b/tests/fixtures/vectorial_parameter_nodes.py @@ -0,0 +1,76 @@ +"""Helpers for tests that copy and pickle vectorial parameter nodes.""" + +import copy +import os +import pickle + +import numpy as np + +from policyengine_core.enums import Enum +from policyengine_core.parameters import ParameterNode + +FANCY_INDEXING_DIRECTORY = os.path.join( + os.path.dirname(os.path.dirname(os.path.abspath(__file__))), + "core", + "parameters_fancy_indexing", +) +INSTANT = "2015-01-01" + +# ``rate...`` in that directory's rate.yaml, at INSTANT. +FAMILIES = ("single", "couple") +HOUSINGS = ("owner", "tenant") +ZONES = ("z1", "z2") +RATE = { + "single": { + "owner": {"z1": 100, "z2": 200}, + "tenant": {"z1": 300, "z2": 400}, + }, + "couple": { + "owner": {"z1": 500, "z2": 600}, + "tenant": {"z1": 700, "z2": 800}, + }, +} + + +class Family(Enum): + single = "Single" + couple = "Couple" + + +class Housing(Enum): + owner = "Owner" + tenant = "Tenant" + + +class SlottedRecords(np.recarray): + """A vector whose class defines ``__slots__``. + + Pickle protocols 0 and 1 refuse an object whose ``__slots__`` they can + read unless its class defines ``__getstate__``. + """ + + __slots__ = ("note",) + + +def fancy_indexing_parameters() -> ParameterNode: + """A new tree, so that no test reads another's cached nodes.""" + return ParameterNode(directory_path=FANCY_INDEXING_DIRECTORY) + + +def _pickle_round_trip(protocol): + def round_trip(obj): + return pickle.loads(pickle.dumps(obj, protocol=protocol)) + + return round_trip + + +COPIERS = { + "copy": copy.copy, + "deepcopy": copy.deepcopy, + **{ + f"pickle-{protocol}": _pickle_round_trip(protocol) + for protocol in range(pickle.HIGHEST_PROTOCOL + 1) + }, +} +# The copiers whose copy has a vector of its own. +DEEP_COPIERS = {name: copier for name, copier in COPIERS.items() if name != "copy"}