Let simulations be deep-copied and unpickled

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.

policyengine_core/enums/enum_array.py

@@ -1,14 +1,47 @@
 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
 
 if typing.TYPE_CHECKING:
     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.

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(