Let vectorial parameter nodes be copied and pickled

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`.

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):

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)