From 404a3a382d1449d94d71e42b476b76f3ab4b1ddd Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Fri, 2 Oct 2026 08:35:41 -0400 Subject: [PATCH 1/4] Resolve ParameterNodeAtInstant children on first read Building every descendant in the constructor made each instant a system is asked about cost a copy of the whole parameter tree, kept for the system's lifetime. For policyengine-us that is about 34,000 node objects, 16-33 MB and several seconds per instant, per tax-benefit system. Co-Authored-By: Claude Opus 5.5 --- .../lazy-parameter-node-at-instant.changed.md | 1 + .../parameters/parameter_node_at_instant.py | 81 +++++++- .../test_parameter_node_at_instant_lazy.py | 194 ++++++++++++++++++ ...test_parameter_node_at_instant_property.py | 91 ++++++++ tests/fixtures/parameter_nodes_at_instant.py | 119 +++++++++++ 5 files changed, 478 insertions(+), 8 deletions(-) create mode 100644 changelog.d/lazy-parameter-node-at-instant.changed.md create mode 100644 tests/core/parameters/test_parameter_node_at_instant_lazy.py create mode 100644 tests/core/parameters/test_parameter_node_at_instant_property.py create mode 100644 tests/fixtures/parameter_nodes_at_instant.py diff --git a/changelog.d/lazy-parameter-node-at-instant.changed.md b/changelog.d/lazy-parameter-node-at-instant.changed.md new file mode 100644 index 000000000..00e96589b --- /dev/null +++ b/changelog.d/lazy-parameter-node-at-instant.changed.md @@ -0,0 +1 @@ +Resolve `ParameterNodeAtInstant` children on first read instead of building the whole parameter tree for every instant, which cost policyengine-us 16 to 33 MB and several seconds per instant on every tax-benefit system. diff --git a/policyengine_core/parameters/parameter_node_at_instant.py b/policyengine_core/parameters/parameter_node_at_instant.py index b55889eb6..776550569 100644 --- a/policyengine_core/parameters/parameter_node_at_instant.py +++ b/policyengine_core/parameters/parameter_node_at_instant.py @@ -16,9 +16,23 @@ ) +# Attributes of the node itself. A lookup of one that is not set is an +# ordinary missing attribute, never a parameter. +_TECHNICAL_ATTRIBUTES = frozenset( + {"_name", "_instant_str", "_node", "_resolved", "_children", "_vectorial_node"} +) + + class ParameterNodeAtInstant: """ Parameter node of the legislation, at a given instant. + + Children resolve on first read and are kept from then on. Building every + descendant up front made each instant a model is asked about cost a copy of + the whole parameter tree (16 to 33 MB and several seconds per instant for + policyengine-us), which every tax-benefit system then kept for its + lifetime. A child that has not been read yet takes the value its parameter + has when it is read. """ def __init__(self, name: str, node: "ParameterNode", instant_str: str): @@ -31,23 +45,68 @@ def __init__(self, name: str, node: "ParameterNode", instant_str: str): # The "technical" attributes are hidden, so that the node children can be easily browsed with auto-completion without pollution self._name = name self._instant_str = instant_str - self._children = {} + # The node whose children are still to be resolved; ``None`` once + # ``_children`` holds all of them. + self._node = node + # Children resolved so far, by name. Becomes ``_children`` itself once + # every child is resolved. + self._resolved = {} - for child_name, child in node.children.items(): - child_at_instant = child._get_at_instant(instant_str) + def _resolve(self, child_name: str): + """Resolve one child, returning ``None`` if this node has no such + child at its instant.""" + node = self._node + child = None if node is None else node.children.get(child_name) + if child is None: + return None + child_at_instant = child._get_at_instant(self._instant_str) + if child_at_instant is not None: + self.add_child(child_name, child_at_instant) + return child_at_instant + + def _resolve_all(self) -> dict: + """Resolve every child, and return them in the parameter node's order.""" + resolved = self._resolved + children = {} + for child_name, child in self._node.children.items(): + if child_name in resolved: + children[child_name] = resolved[child_name] + continue + child_at_instant = child._get_at_instant(self._instant_str) if child_at_instant is not None: - self.add_child(child_name, child_at_instant) + children[child_name] = child_at_instant + if child_name not in _TECHNICAL_ATTRIBUTES: + setattr(self, child_name, child_at_instant) + for child_name, child_at_instant in resolved.items(): + # Children added with ``add_child`` that the node does not have. + children.setdefault(child_name, child_at_instant) + self._children = self._resolved = children + self._node = None + return children def add_child(self, child_name: str, child_at_instant: "ParameterNodeAtInstant"): - self._children[child_name] = child_at_instant - setattr(self, child_name, child_at_instant) + self._resolved[child_name] = child_at_instant + if child_name not in _TECHNICAL_ATTRIBUTES: + setattr(self, child_name, child_at_instant) def __getattr__(self, key: str): - if key == "_name": + # Reached only for a name that is not an instance attribute yet: a + # child that has not been read, ``_children`` before every child has + # been resolved, or a name this node does not have. + if key in _TECHNICAL_ATTRIBUTES: + if key == "_children" and self.__dict__.get("_node") is not None: + return self._resolve_all() raise AttributeError(key) + if self.__dict__.get("_node") is not None: + child_at_instant = self._resolve(key) + if child_at_instant is not None: + return child_at_instant param_name = helpers._compose_name(self._name, item_name=key) raise ParameterNotFoundError(param_name, self._instant_str) + def __dir__(self): + return sorted(set(super().__dir__()) | set(map(str, self._children))) + def __getitem__( self, key: str ) -> Union["ParameterNodeAtInstant", VectorialParameterNodeAtInstant]: @@ -68,7 +127,13 @@ def __getitem__( ) self._vectorial_node = vectorial return vectorial[key] - return self._children[key] + resolved = self._resolved + if key in resolved: + return resolved[key] + child_at_instant = self._resolve(key) + if child_at_instant is None: + raise KeyError(key) + return child_at_instant def __iter__(self) -> Iterable: return iter(self._children) diff --git a/tests/core/parameters/test_parameter_node_at_instant_lazy.py b/tests/core/parameters/test_parameter_node_at_instant_lazy.py new file mode 100644 index 000000000..81fbf2e1d --- /dev/null +++ b/tests/core/parameters/test_parameter_node_at_instant_lazy.py @@ -0,0 +1,194 @@ +"""``ParameterNodeAtInstant`` resolves its children on first read. + +It used to build every descendant in its constructor, so each instant a +system was asked about cost a copy of the whole parameter tree (about 34,000 +node objects and 16-33 MB per instant for policyengine-us) that the system +then kept. These tests pin that nothing is built before it is read, and that +everything read is what the up-front build held. +""" + +import copy +import gc + +import numpy as np +import pytest + +from policyengine_core.country_template import CountryTaxBenefitSystem +from policyengine_core.errors import ParameterNotFoundError +from policyengine_core.parameters import ParameterNodeAtInstant +from tests.fixtures.parameter_nodes_at_instant import ( + INSTANTS, + build_tree, + materialise, + ordered, + snapshot, +) + + +@pytest.mark.parametrize("instant", INSTANTS) +def test_fully_read_node_matches_the_up_front_build(instant): + tree = build_tree() + assert ordered(materialise(tree.get_at_instant(instant))) == ordered( + snapshot(tree, instant) + ) + + +@pytest.mark.parametrize( + "instant", ["2000-01-01", "2012-06-01", "2015-01-01", "2017-01-01", "2030-01-01"] +) +def test_country_template_matches_the_up_front_build(instant): + parameters = CountryTaxBenefitSystem().parameters + assert ordered(materialise(parameters.get_at_instant(instant))) == ordered( + snapshot(parameters, instant) + ) + + +def test_nothing_is_resolved_until_it_is_read(): + at_instant = build_tree().get_at_instant("2017-03-01") + assert at_instant._resolved == {} + + assert at_instant.group.x == 0.1 + # Only the path that was read: the root's ``group`` and the group's ``x``. + assert list(at_instant._resolved) == ["group"] + assert list(at_instant.group._resolved) == ["x"] + + +def test_asking_for_many_instants_builds_one_node_each(): + tree = build_tree() + gc.collect() + before = sum(isinstance(o, ParameterNodeAtInstant) for o in gc.get_objects()) + for year in range(2000, 2100): + tree.get_at_instant(f"{year}-01-01") + after = sum(isinstance(o, ParameterNodeAtInstant) for o in gc.get_objects()) + # The up-front build made 9 nodes per instant for this tree (900 here). + assert after - before == 100 + + +def test_a_read_child_is_resolved_once(): + at_instant = build_tree().get_at_instant("2017-03-01") + # A scale is rebuilt by every ``_get_at_instant`` call, so identity shows + # that the node kept what it resolved. + assert at_instant.scale is at_instant.scale + assert at_instant["scale"] is at_instant.scale + assert at_instant._children["scale"] is at_instant.scale + assert at_instant.group is at_instant["group"] + + +def test_children_are_in_parameter_order_whatever_was_read_first(): + tree = build_tree() + at_instant = tree.get_at_instant("2030-01-01") + at_instant.by_zone + at_instant["late"] + at_instant.group + assert list(at_instant) == list(snapshot(tree, "2030-01-01")) + assert list(at_instant._children) == list(snapshot(tree, "2030-01-01")) + + +def test_child_without_a_value_yet_is_absent(): + at_instant = build_tree().get_at_instant("2017-03-01") + with pytest.raises(ParameterNotFoundError): + at_instant.late + with pytest.raises(ParameterNotFoundError) as error: + at_instant.group.y + assert str(error.value) == ( + "The parameter 'group[y]' was not found in the 2017-03-01 " + "tax and benefit system." + ) + with pytest.raises(KeyError): + at_instant["late"] + assert not hasattr(at_instant, "late") + assert "late" not in at_instant + assert "late" not in at_instant._children + assert at_instant.group.inner._children == {} + + +def test_unknown_names_raise_as_before(): + at_instant = build_tree().get_at_instant("2017-03-01") + for read_first in (False, True): + if read_first: + at_instant._children + with pytest.raises(ParameterNotFoundError) as error: + at_instant.group.nope + assert str(error.value) == ( + "The parameter 'group[nope]' was not found in the 2017-03-01 " + "tax and benefit system." + ) + with pytest.raises(KeyError): + at_instant.group["nope"] + assert not hasattr(at_instant, "nope") + + +def test_vectorial_indexing_reads_the_whole_subtree(): + at_instant = build_tree().get_at_instant("2020-01-01") + zones = np.array(["z1", "z3", "z2"]) + np.testing.assert_array_equal(at_instant.by_zone[zones].amount, [200, 600, 400]) + np.testing.assert_allclose(at_instant.by_zone[zones].rate, [0.1, 0.3, 0.2]) + # Nothing outside the indexed subtree was resolved for it. + assert list(at_instant._resolved) == ["by_zone"] + + +def test_dir_lists_children_for_completion(): + at_instant = build_tree().get_at_instant("2017-03-01") + names = dir(at_instant) + assert {"flat_rate", "group", "scale", "by_zone"} <= set(names) + assert "late" not in names + + +def test_repr_shows_every_child(): + text = repr(build_tree().get_at_instant("2017-03-01")) + for name in ("flat_rate", "group", "scale", "by_zone", "z3"): + assert name in text + + +def test_add_child_before_and_after_a_full_read(): + at_instant = build_tree().get_at_instant("2017-03-01") + at_instant.add_child("extra", 1) + assert at_instant.extra == 1 and at_instant["extra"] == 1 + assert list(at_instant._children)[-1] == "extra" + at_instant.add_child("later", 2) + assert at_instant.later == 2 and at_instant._children["later"] == 2 + + +def test_a_node_asked_for_after_an_update_reads_the_update(): + tree = build_tree() + assert tree.get_at_instant("2017-03-01").flat_rate == 0.1 + tree.flat_rate.update(value=0.5, period="year:2017:1") + assert tree.get_at_instant("2017-03-01").flat_rate == 0.5 + + +def test_a_child_read_before_an_update_keeps_what_was_read(): + tree = build_tree() + held = tree.get_at_instant("2017-03-01") + assert held.flat_rate == 0.1 + tree.flat_rate.update(value=0.5, period="year:2017:1") + assert held.flat_rate == 0.1 + + +def test_a_child_not_read_before_an_update_reads_the_update(): + """Intended difference from the up-front build, which held 0.1 here: a + node held across a change to its parameters reads, for each child, the + value at the time that child is first read.""" + tree = build_tree() + held = tree.get_at_instant("2017-03-01") + tree.flat_rate.update(value=0.5, period="year:2017:1") + assert held.flat_rate == 0.5 + + +def test_deep_copied_tree_keeps_its_own_nodes_at_instant(): + tree = build_tree() + tree.get_at_instant("2017-03-01").group.x # cache a partly read node + clone = copy.deepcopy(tree) + cached = clone._at_instant_cache["2017-03-01"] + clone.group.y.update(value=0.9, period="year:2010:30") + clone.flat_rate.update(value=0.7, period="year:2010:30") + # The copy's cached node reads the copy, not the original tree. + assert cached.flat_rate == 0.7 + assert tree.get_at_instant("2017-03-01").flat_rate == 0.1 + assert not hasattr(tree.get_at_instant("2017-03-01").group, "y") + + +def test_copy_of_a_node_at_instant_reads_the_same_values(): + at_instant = build_tree().get_at_instant("2017-03-01") + at_instant.group + for duplicate in (copy.copy(at_instant), copy.deepcopy(at_instant)): + assert materialise(duplicate) == materialise(at_instant) diff --git a/tests/core/parameters/test_parameter_node_at_instant_property.py b/tests/core/parameters/test_parameter_node_at_instant_property.py new file mode 100644 index 000000000..5eeec9372 --- /dev/null +++ b/tests/core/parameters/test_parameter_node_at_instant_property.py @@ -0,0 +1,91 @@ +"""Reads of a lazily resolved ``ParameterNodeAtInstant`` in any order. + +For any instant and any sequence of reads (attribute, item, iteration, +``_children``, ``dir``, ``repr``, vectorial indexing, names that do not +exist), every read returns what the up-front build held, and the node read in +full afterwards equals that build. ``test_parameter_node_at_instant_lazy.py`` +pins the same behaviour with examples. +""" + +from __future__ import annotations + +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.errors import ParameterNotFoundError +from policyengine_core.parameters import ParameterNodeAtInstant +from tests.fixtures.parameter_nodes_at_instant import ( + INSTANTS, + build_tree, + describe, + materialise, + ordered, + snapshot, +) + +# Node paths in the fixture tree, and names to ask each for (present at some +# instants, absent at others, or never there). +PATHS = [(), ("group",), ("group", "inner"), ("by_zone",), ("by_zone", "z2")] +NAMES = ["flat_rate", "late", "group", "scale", "by_zone", "x", "y", "inner"] +NAMES += ["1", "2", "z1", "z2", "z3", "amount", "rate", "nope"] +KINDS = ["getattr", "getitem", "iter", "children", "dir", "repr", "vector"] + +reads = st.tuples( + st.sampled_from(PATHS), st.sampled_from(KINDS), st.sampled_from(NAMES) +) + + +def _comparable(value): + if isinstance(value, ParameterNodeAtInstant): + return materialise(value) + return describe(value) + + +@hypothesis.given( + instant=st.sampled_from(INSTANTS), sequence=st.lists(reads, max_size=30) +) +@hypothesis.settings(max_examples=400, deadline=None) +def test_reads_in_any_order_match_the_up_front_build(instant, sequence): + tree = build_tree() + expected_root = snapshot(tree, instant) + root = tree.get_at_instant(instant) + + for path, kind, name in sequence: + node, expected = root, expected_root + for step in path: + node, expected = node[step], expected[step] + + if kind == "getattr": + if name in expected: + assert _comparable(getattr(node, name)) == expected[name] + else: + with pytest.raises(ParameterNotFoundError): + getattr(node, name) + assert not hasattr(node, name) + elif kind == "getitem": + if name in expected: + assert _comparable(node[name]) == expected[name] + else: + with pytest.raises(KeyError): + node[name] + elif kind == "iter": + assert list(node) == list(expected) + assert (name in node) == (name in expected) + elif kind == "children": + assert ordered(materialise(node)) == ordered(expected) + elif kind == "dir": + assert set(expected) <= set(dir(node)) + elif kind == "repr": + assert all(child in repr(node) for child in expected) + elif kind == "vector" and path == ("by_zone",) and expected["z1"]: + zones = np.array(["z3", "z1"]) + np.testing.assert_array_equal( + node[zones].amount, + [expected["z3"]["amount"], expected["z1"]["amount"]], + ) + + assert ordered(materialise(root)) == ordered(expected_root) diff --git a/tests/fixtures/parameter_nodes_at_instant.py b/tests/fixtures/parameter_nodes_at_instant.py new file mode 100644 index 000000000..70198ed31 --- /dev/null +++ b/tests/fixtures/parameter_nodes_at_instant.py @@ -0,0 +1,119 @@ +"""A parameter tree and an up-front reference for the at-instant tests. + +``snapshot`` computes what a ``ParameterNodeAtInstant`` held when it built +every descendant in its constructor, straight from the parameter tree and +without using ``ParameterNodeAtInstant``. ``materialise`` reads the same +thing out of a ``ParameterNodeAtInstant``. +""" + +from policyengine_core.parameters import ( + Parameter, + ParameterNode, + ParameterNodeAtInstant, + ParameterScale, +) + +INSTANTS = [ + "2009-12-31", + "2010-01-01", + "2014-06-30", + "2015-01-01", + "2017-03-01", + "2018-01-01", + "2020-01-01", + "2029-12-31", + "2030-01-01", + "2035-07-01", +] + + +def _zone(amount, rate): + return { + "amount": {"values": {"2010-01-01": amount, "2020-01-01": amount * 2}}, + "rate": {"values": {"2010-01-01": rate}}, + } + + +TREE = { + "flat_rate": {"values": {"2015-01-01": 0.1, "2020-01-01": 0.2}}, + # Not defined before 2030, so absent from earlier instants. + "late": {"values": {"2030-01-01": 5}}, + "group": { + "x": {"values": {"2010-01-01": 0.1}}, + "y": {"values": {"2018-01-01": 0.2}}, + # Every child undefined before 2018: an empty node, not an absent one. + "inner": { + "1": {"values": {"2018-01-01": 10}}, + "2": {"values": {"2018-01-01": 20, "2020-01-01": 30}}, + }, + }, + "scale": { + "brackets": [ + { + "threshold": {"2010-01-01": {"value": 0}}, + "rate": {"2010-01-01": {"value": 0.1}}, + }, + { + "threshold": {"2010-01-01": {"value": 1000}}, + "rate": { + "2010-01-01": {"value": 0.2}, + "2020-01-01": {"value": 0.3}, + }, + }, + ] + }, + # Homogeneous children: vectorisable with ``node[array]``. + "by_zone": {"z1": _zone(100, 0.1), "z2": _zone(200, 0.2), "z3": _zone(300, 0.3)}, +} + + +def build_tree() -> ParameterNode: + return ParameterNode("", data=TREE) + + +def describe(value): + """A comparable form of one at-instant value.""" + if isinstance(value, (dict, int, float, bool, str)) or value is None: + return value + # Tax scales refuse ``==``. + return ( + type(value).__name__, + tuple(value.thresholds), + tuple(getattr(value, "rates", getattr(value, "amounts", ()))), + ) + + +def snapshot(node: ParameterNode, instant_str: str) -> dict: + """Every descendant's value at the instant, in the node's own order, + leaving out children that have no value yet.""" + result = {} + for name, child in node.children.items(): + if isinstance(child, ParameterScale): + value = describe(child._get_at_instant(instant_str)) + elif isinstance(child, Parameter): + value = child._get_at_instant(instant_str) + else: + value = snapshot(child, instant_str) + if value is not None: + result[name] = value + return result + + +def materialise(node_at_instant: ParameterNodeAtInstant) -> dict: + """The same nested mapping, read from a ``ParameterNodeAtInstant``.""" + return { + name: ( + materialise(child) + if isinstance(child, ParameterNodeAtInstant) + else describe(child) + ) + for name, child in node_at_instant._children.items() + } + + +def ordered(mapping: dict): + """A form that compares key order as well as content.""" + return [ + (key, ordered(value) if isinstance(value, dict) else value) + for key, value in mapping.items() + ] From a90b524a5d47a654e17372c4bee084f19e06ccb5 Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Fri, 2 Oct 2026 09:52:14 -0400 Subject: [PATCH 2/4] Make lazy at-instant resolution thread-safe, collision-proof and copy-coherent Review of 404a3a38 found five regressions against the up-front build: - concurrent first reads could raise KeyError/AttributeError or resolve a child twice: first resolution now runs under a module RLock; - children named like the new state or helpers (_node, _resolved, _resolve, ...) read as internal objects: state is name-mangled, helpers are module functions, and children named like the node's own attributes resolve at construction, as before; - a shallow copy and its original could resolve a scale twice: both now read one grow-only resolved map; - a child with no value at the instant was re-resolved on every read: its absence is cached; - deepcopy/pickle of a traced tree's raw node recursed through TracingParameterNodeAtInstant.__getattr__, which now raises AttributeError for its own unset fields. A node also drops its parameter node once every child has been read. Co-Authored-By: Claude Opus 5.5 --- .../parameters/parameter_node_at_instant.py | 198 ++++++++++++------ .../tracing_parameter_node_at_instant.py | 6 + .../test_parameter_node_at_instant_lazy.py | 181 +++++++++++++++- 3 files changed, 316 insertions(+), 69 deletions(-) diff --git a/policyengine_core/parameters/parameter_node_at_instant.py b/policyengine_core/parameters/parameter_node_at_instant.py index 776550569..010d5f545 100644 --- a/policyengine_core/parameters/parameter_node_at_instant.py +++ b/policyengine_core/parameters/parameter_node_at_instant.py @@ -1,5 +1,6 @@ import os import sys +import threading from typing import TYPE_CHECKING, Iterable, Union import numpy @@ -16,11 +17,20 @@ ) -# Attributes of the node itself. A lookup of one that is not set is an -# ordinary missing attribute, never a parameter. -_TECHNICAL_ATTRIBUTES = frozenset( - {"_name", "_instant_str", "_node", "_resolved", "_children", "_vectorial_node"} -) +# A node's own state, under name-mangled keys so that no parameter child can +# shadow it. +_NODE = "_ParameterNodeAtInstant__node" +_INSTANT = "_ParameterNodeAtInstant__instant" +_RESOLVED = "_ParameterNodeAtInstant__resolved" +_ABSENT = "_ParameterNodeAtInstant__absent" +_CHILDREN = "_ParameterNodeAtInstant__children" +_MISSING = object() + +# Held while a child is resolved for the first time, so that concurrent +# readers never see a half-resolved node or resolve one child twice. Reads of +# a resolved child never take it. Reentrant because resolving a child can +# build that child's own node. +_RESOLUTION_LOCK = threading.RLock() class ParameterNodeAtInstant: @@ -32,7 +42,8 @@ class ParameterNodeAtInstant: the whole parameter tree (16 to 33 MB and several seconds per instant for policyengine-us), which every tax-benefit system then kept for its lifetime. A child that has not been read yet takes the value its parameter - has when it is read. + has when it is first read; a child read once, including one found to have + no value at this instant, keeps what that read found. """ def __init__(self, name: str, node: "ParameterNode", instant_str: str): @@ -45,67 +56,51 @@ def __init__(self, name: str, node: "ParameterNode", instant_str: str): # The "technical" attributes are hidden, so that the node children can be easily browsed with auto-completion without pollution self._name = name self._instant_str = instant_str - # The node whose children are still to be resolved; ``None`` once - # ``_children`` holds all of them. - self._node = node - # Children resolved so far, by name. Becomes ``_children`` itself once - # every child is resolved. - self._resolved = {} - - def _resolve(self, child_name: str): - """Resolve one child, returning ``None`` if this node has no such - child at its instant.""" - node = self._node - child = None if node is None else node.children.get(child_name) - if child is None: - return None - child_at_instant = child._get_at_instant(self._instant_str) - if child_at_instant is not None: - self.add_child(child_name, child_at_instant) - return child_at_instant - - def _resolve_all(self) -> dict: - """Resolve every child, and return them in the parameter node's order.""" - resolved = self._resolved - children = {} - for child_name, child in self._node.children.items(): - if child_name in resolved: - children[child_name] = resolved[child_name] - continue - child_at_instant = child._get_at_instant(self._instant_str) - if child_at_instant is not None: - children[child_name] = child_at_instant - if child_name not in _TECHNICAL_ATTRIBUTES: - setattr(self, child_name, child_at_instant) - for child_name, child_at_instant in resolved.items(): - # Children added with ``add_child`` that the node does not have. - children.setdefault(child_name, child_at_instant) - self._children = self._resolved = children - self._node = None - return children + state = self.__dict__ + # The parameter node children still resolve from; None once every + # child is resolved. + state[_NODE] = node + state[_INSTANT] = instant_str + # Children resolved so far. Only ever grows, and a shallow copy shares + # it, so every copy reads the same child objects. + state[_RESOLVED] = {} + # Children of the node that have no value at this instant. + state[_ABSENT] = set() + # A child whose name the node itself uses (``add_child``, ``_name``, + # ...) shadows that name, as when every child was set here. + for child_name in _SHADOWING_NAMES: + if child_name in node.children: + _resolve(self, child_name) def add_child(self, child_name: str, child_at_instant: "ParameterNodeAtInstant"): - self._resolved[child_name] = child_at_instant - if child_name not in _TECHNICAL_ATTRIBUTES: - setattr(self, child_name, child_at_instant) + with _RESOLUTION_LOCK: + _store(self, child_name, child_at_instant) def __getattr__(self, key: str): - # Reached only for a name that is not an instance attribute yet: a - # child that has not been read, ``_children`` before every child has - # been resolved, or a name this node does not have. - if key in _TECHNICAL_ATTRIBUTES: - if key == "_children" and self.__dict__.get("_node") is not None: - return self._resolve_all() + # Reached only for a name that is not an instance attribute: a child + # not read yet, ``_children`` before every child is resolved, or a name + # this node does not have. + if key.startswith("__") and key.endswith("__"): + # Protocol lookups (copy, pickle, numpy) are never parameters. + raise AttributeError(key) + if _RESOLVED not in self.__dict__ or key == "_name": + # An instance still being rebuilt by copy or pickle. raise AttributeError(key) - if self.__dict__.get("_node") is not None: - child_at_instant = self._resolve(key) - if child_at_instant is not None: - return child_at_instant + if key == "_children": + return _materialise(self) + child_at_instant = _resolve(self, key) + if child_at_instant is not _MISSING: + return child_at_instant param_name = helpers._compose_name(self._name, item_name=key) raise ParameterNotFoundError(param_name, self._instant_str) def __dir__(self): - return sorted(set(super().__dir__()) | set(map(str, self._children))) + self._children + return [ + name + for name in super().__dir__() + if not name.startswith("_ParameterNodeAtInstant__") + ] def __getitem__( self, key: str @@ -127,12 +122,11 @@ def __getitem__( ) self._vectorial_node = vectorial return vectorial[key] - resolved = self._resolved - if key in resolved: - return resolved[key] - child_at_instant = self._resolve(key) - if child_at_instant is None: - raise KeyError(key) + child_at_instant = self.__dict__[_RESOLVED].get(key, _MISSING) + if child_at_instant is _MISSING: + child_at_instant = _resolve(self, key) + if child_at_instant is _MISSING: + raise KeyError(key) return child_at_instant def __iter__(self) -> Iterable: @@ -146,3 +140,79 @@ def __repr__(self) -> str: ] ) return result + + +def _store(node_at_instant, child_name, child_at_instant): + """Record a child. The caller holds ``_RESOLUTION_LOCK``.""" + state = node_at_instant.__dict__ + state[_RESOLVED][child_name] = child_at_instant + state[_ABSENT].discard(child_name) + children = state.get(_CHILDREN) + if children is not None: + children[child_name] = child_at_instant + setattr(node_at_instant, child_name, child_at_instant) + + +def _resolve(node_at_instant, child_name): + """Return one child, resolving it if no read has yet; ``_MISSING`` if the + node has no such child or it has no value at this instant.""" + state = node_at_instant.__dict__ + with _RESOLUTION_LOCK: + resolved = state[_RESOLVED] + if child_name in resolved: + return resolved[child_name] + if child_name in state[_ABSENT]: + return _MISSING + node = state[_NODE] + child = None if node is None else node.children.get(child_name) + if child is None: + return _MISSING + child_at_instant = child._get_at_instant(state[_INSTANT]) + if child_at_instant is None: + state[_ABSENT].add(child_name) + else: + _store(node_at_instant, child_name, child_at_instant) + if len(resolved) + len(state[_ABSENT]) >= len(node.children): + # Every child has been read: drop the parameter node. + _materialise(node_at_instant) + return _MISSING if child_at_instant is None else child_at_instant + + +def _materialise(node_at_instant): + """Resolve every child; return them in the parameter node's order, with + any added by ``add_child`` after them.""" + state = node_at_instant.__dict__ + with _RESOLUTION_LOCK: + children = state.get(_CHILDREN) + if children is not None: + return children + node = state[_NODE] + children = {} + for child_name in node.children if node is not None else (): + child_at_instant = _resolve(node_at_instant, child_name) + if child_at_instant is not _MISSING: + children[child_name] = child_at_instant + if state.get(_CHILDREN) is not None: + # Resolving the last child completed the node already. + return state[_CHILDREN] + for child_name, child_at_instant in state[_RESOLVED].items(): + children.setdefault(child_name, child_at_instant) + state[_CHILDREN] = children + if "_children" not in state: + # Unless a child named ``_children`` took the name. + state["_children"] = children + # Every child is resolved: the parameter node is no longer needed. + state[_NODE] = None + return children + + +# Names a ParameterNodeAtInstant itself answers to, which a child of the same +# name shadows. +_SHADOWING_NAMES = frozenset( + {"_name", "_instant_str", "_children", "_vectorial_node"} + | { + name + for name in dir(ParameterNodeAtInstant) + if not name.startswith("__") and not name.startswith("_ParameterNodeAtInstant") + } +) diff --git a/policyengine_core/tracers/tracing_parameter_node_at_instant.py b/policyengine_core/tracers/tracing_parameter_node_at_instant.py index d4bf17466..394e1c540 100644 --- a/policyengine_core/tracers/tracing_parameter_node_at_instant.py +++ b/policyengine_core/tracers/tracing_parameter_node_at_instant.py @@ -21,6 +21,9 @@ Child = Union[ParameterNode, ArrayLike] +_OWN_ATTRIBUTES = frozenset({"parameter_node_at_instant", "tracer", "branch_name"}) + + class TracingParameterNodeAtInstant: def __init__( self, @@ -36,6 +39,9 @@ def __getattr__( self, key: str, ) -> Union[TracingParameterNodeAtInstant, Child]: + if key in _OWN_ATTRIBUTES: + # Not set yet: copy or pickle is rebuilding this wrapper. + raise AttributeError(key) child = getattr(self.parameter_node_at_instant, key) return self.get_traced_child(child, key) diff --git a/tests/core/parameters/test_parameter_node_at_instant_lazy.py b/tests/core/parameters/test_parameter_node_at_instant_lazy.py index 81fbf2e1d..ff2604e02 100644 --- a/tests/core/parameters/test_parameter_node_at_instant_lazy.py +++ b/tests/core/parameters/test_parameter_node_at_instant_lazy.py @@ -9,13 +9,18 @@ import copy import gc +import pickle +import threading +import weakref import numpy as np import pytest from policyengine_core.country_template import CountryTaxBenefitSystem from policyengine_core.errors import ParameterNotFoundError -from policyengine_core.parameters import ParameterNodeAtInstant +from policyengine_core.parameters import ParameterNode, ParameterNodeAtInstant +from policyengine_core.parameters.parameter_node_at_instant import _RESOLVED +from policyengine_core.tracers import FullTracer from tests.fixtures.parameter_nodes_at_instant import ( INSTANTS, build_tree, @@ -43,14 +48,18 @@ def test_country_template_matches_the_up_front_build(instant): ) +def resolved(node_at_instant): + return list(node_at_instant.__dict__[_RESOLVED]) + + def test_nothing_is_resolved_until_it_is_read(): at_instant = build_tree().get_at_instant("2017-03-01") - assert at_instant._resolved == {} + assert resolved(at_instant) == [] assert at_instant.group.x == 0.1 # Only the path that was read: the root's ``group`` and the group's ``x``. - assert list(at_instant._resolved) == ["group"] - assert list(at_instant.group._resolved) == ["x"] + assert resolved(at_instant) == ["group"] + assert resolved(at_instant.group) == ["x"] def test_asking_for_many_instants_builds_one_node_each(): @@ -124,7 +133,7 @@ def test_vectorial_indexing_reads_the_whole_subtree(): np.testing.assert_array_equal(at_instant.by_zone[zones].amount, [200, 600, 400]) np.testing.assert_allclose(at_instant.by_zone[zones].rate, [0.1, 0.3, 0.2]) # Nothing outside the indexed subtree was resolved for it. - assert list(at_instant._resolved) == ["by_zone"] + assert resolved(at_instant) == ["by_zone"] def test_dir_lists_children_for_completion(): @@ -192,3 +201,165 @@ def test_copy_of_a_node_at_instant_reads_the_same_values(): at_instant.group for duplicate in (copy.copy(at_instant), copy.deepcopy(at_instant)): assert materialise(duplicate) == materialise(at_instant) + + +@pytest.mark.parametrize( + "child_name", + # Not ``add_child``: ParameterNode itself cannot hold a child of that name. + ["_name", "_instant_str", "_vectorial_node", "_node", "_resolved"] + + ["_resolve", "_resolve_all", "_store", "_materialise"], +) +def test_children_named_like_the_nodes_own_attributes(child_name): + """Such a child reads as its value whatever the read order, as when every + child was set up front, and leaves its siblings readable.""" + tree = ParameterNode( + "", + data={ + child_name: {"values": {"2010-01-01": 42}}, + "ok": {"values": {"2010-01-01": 7}}, + }, + ) + for read_name_first in (True, False): + tree._at_instant_cache.clear() + at_instant = tree.get_at_instant("2017-03-01") + if read_name_first: + assert getattr(at_instant, child_name) == 42 + assert at_instant.ok == 7 and at_instant["ok"] == 7 + else: + assert at_instant["ok"] == 7 and at_instant.ok == 7 + assert getattr(at_instant, child_name) == 42 + assert at_instant[child_name] == 42 + assert set(at_instant._children) == {child_name, "ok"} + + +class BlockingChild: + """A parameter child whose first resolution waits until released.""" + + def __init__(self): + self.calls = 0 + self.entered = threading.Event() + self.release = threading.Event() + + def _get_at_instant(self, instant): + self.calls += 1 + if self.calls == 1: + self.entered.set() + assert self.release.wait(10), "never released" + return object() + + +@pytest.mark.parametrize( + "second_read", + [ + lambda node: node.slow, + lambda node: node["slow"], + lambda node: node._children["slow"], + lambda node: [node[name] for name in node][0], + ], +) +def test_concurrent_first_reads_resolve_a_child_once(second_read): + tree = ParameterNode("", data={"other": {"values": {"2010-01-01": 1}}}) + slow = BlockingChild() + tree.children = {"slow": slow, **tree.children} + at_instant = tree.get_at_instant("2017-03-01") + results, errors = [], [] + + def read(reader): + try: + results.append(reader(at_instant)) + except Exception as error: # noqa: BLE001 - recorded for the assertion + errors.append(error) + + first = threading.Thread(target=read, args=(lambda node: node.slow,)) + first.start() + assert slow.entered.wait(10) + second = threading.Thread(target=read, args=(second_read,)) + second.start() + second.join(0.2) # let the second reader reach the child too + slow.release.set() + first.join(10) + second.join(10) + + assert not errors, errors + assert len(results) == 2 and results[0] is results[1] + assert slow.calls == 1 + + +def test_shallow_copies_read_the_same_children(): + held = build_tree().get_at_instant("2017-03-01") + duplicate = copy.copy(held) + original_scale = held.scale + # A scale is rebuilt by each resolution, so identity shows one resolution. + assert duplicate.scale is original_scale + assert duplicate["scale"] is original_scale + assert held["scale"] is original_scale + duplicate._children + assert held._children["scale"] is original_scale + assert held.by_zone is duplicate["by_zone"] + + +@pytest.mark.parametrize( + "duplicate", [copy.deepcopy, lambda n: pickle.loads(pickle.dumps(n))] +) +def test_traced_nodes_copy_and_pickle(duplicate): + tree = build_tree() + tree.trace = True + tree.tracer = FullTracer() + tree.branch_name = "default" + traced = tree.get_at_instant("2017-03-01") + traced.group.x # partly resolved, through the tracing wrapper + raw = traced.parameter_node_at_instant + for node in (raw, traced): + assert duplicate(node).flat_rate == 0.1 + assert duplicate(node).group.x == 0.1 + + +def test_a_child_without_a_value_is_resolved_once(): + tree = ParameterNode("", data={"late": {"values": {"2030-01-01": 1}}}) + held = tree.get_at_instant("2020-01-01") + calls = [] + original = tree.children["late"]._get_at_instant + + def counting(instant): + calls.append(instant) + return original(instant) + + tree.children["late"]._get_at_instant = counting + for _ in range(3): + with pytest.raises(ParameterNotFoundError): + held.late + with pytest.raises(KeyError): + held["late"] + assert list(held) == [] + assert len(calls) == 1 + + # A child already read keeps what that read found, value or none. + tree.children["late"].update(value=2, period="year:2020:1") + with pytest.raises(KeyError): + held["late"] + assert tree.get_at_instant("2020-01-01").late == 2 + + +def test_a_node_lets_go_of_its_parameter_node_once_every_child_is_read(): + tree = ParameterNode( + "", + data={ + "group": { + "x": {"values": {"2010-01-01": 1}}, + "late": {"values": {"2030-01-01": 2}}, + }, + "unused": {"values": {"2010-01-01": 3}}, + }, + ) + group = tree.get_at_instant("2017-03-01").group + tree_ref = weakref.ref(tree) + del tree + gc.collect() + # Partly read: the node still needs the tree to resolve the rest. + assert tree_ref() is not None + assert group.x == 1 + with pytest.raises(ParameterNotFoundError): + group.late + gc.collect() + assert tree_ref() is None + assert group.x == 1 and list(group) == ["x"] From 4e6bc315236540e1efd1ebd90cfbce6d15e79820 Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Fri, 2 Oct 2026 10:00:12 -0400 Subject: [PATCH 3/4] Compare list and array parameter values in the at-instant reference Co-Authored-By: Claude Opus 5.5 --- tests/fixtures/parameter_nodes_at_instant.py | 26 ++++++++++++++------ 1 file changed, 19 insertions(+), 7 deletions(-) diff --git a/tests/fixtures/parameter_nodes_at_instant.py b/tests/fixtures/parameter_nodes_at_instant.py index 70198ed31..7b991041b 100644 --- a/tests/fixtures/parameter_nodes_at_instant.py +++ b/tests/fixtures/parameter_nodes_at_instant.py @@ -6,6 +6,8 @@ thing out of a ``ParameterNodeAtInstant``. """ +import numpy as np + from policyengine_core.parameters import ( Parameter, ParameterNode, @@ -73,14 +75,24 @@ def build_tree() -> ParameterNode: def describe(value): """A comparable form of one at-instant value.""" - if isinstance(value, (dict, int, float, bool, str)) or value is None: + if isinstance(value, (int, float, bool, str)) or value is None: return value - # Tax scales refuse ``==``. - return ( - type(value).__name__, - tuple(value.thresholds), - tuple(getattr(value, "rates", getattr(value, "amounts", ()))), - ) + if isinstance(value, (list, tuple)): + return tuple(describe(item) for item in value) + if isinstance(value, dict): + return {key: describe(item) for key, item in value.items()} + if hasattr(value, "thresholds"): + # Tax scales refuse ``==``. + return ( + type(value).__name__, + tuple(value.thresholds), + tuple(getattr(value, "rates", getattr(value, "amounts", ()))), + ) + if isinstance(value, np.ndarray): + return ("ndarray", str(value.dtype), value.tolist()) + if isinstance(value, np.generic): + return value.item() + return (type(value).__name__, repr(value)) def snapshot(node: ParameterNode, instant_str: str) -> dict: From de47cde04f9c594b48b3d6c699285fc9bb485872 Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Sat, 3 Oct 2026 08:24:11 -0400 Subject: [PATCH 4/4] Keep lazy at-instant state in one shared object; fix round-2 review cases Round-2 review of a90b524a found reads that still differed from the up-front build. A node now keeps all of its own state in one _LazyChildren object under a single reserved name (a parameter of that name is rejected): - which children a node has, and their order, are fixed at construction (a child deleted from the tree later keeps its place); - completion counts only those children, so add_child extras can no longer send the first read into recursive completion; - resolution results and add_child entries are kept apart: items are shared by shallow copies, attributes belong to each copy, as before; - a copy installs what it reads as attributes (no lock on warm reads); - children named like attributes of the node's runtime type (dunders and subclass methods included) resolve at construction; a subclass overriding add_child gets the up-front build through it; - a reentrant read keeps the first recorded child; - nodes pickled by the up-front build load (__setstate__). Co-Authored-By: Claude Opus 5.5 --- .../parameters/parameter_node_at_instant.py | 259 ++++++++++-------- .../test_parameter_node_at_instant_lazy.py | 151 +++++++++- 2 files changed, 299 insertions(+), 111 deletions(-) diff --git a/policyengine_core/parameters/parameter_node_at_instant.py b/policyengine_core/parameters/parameter_node_at_instant.py index 010d5f545..43787a045 100644 --- a/policyengine_core/parameters/parameter_node_at_instant.py +++ b/policyengine_core/parameters/parameter_node_at_instant.py @@ -17,13 +17,11 @@ ) -# A node's own state, under name-mangled keys so that no parameter child can -# shadow it. -_NODE = "_ParameterNodeAtInstant__node" -_INSTANT = "_ParameterNodeAtInstant__instant" -_RESOLVED = "_ParameterNodeAtInstant__resolved" -_ABSENT = "_ParameterNodeAtInstant__absent" -_CHILDREN = "_ParameterNodeAtInstant__children" +# The one instance attribute a node keeps for itself; a parameter of this +# name is rejected. +_STATE = "_ParameterNodeAtInstant__state" +# Attributes a node sets on itself, which a child of the same name shadows. +_OWN_NAMES = frozenset({"_name", "_instant_str", "_children", "_vectorial_node"}) _MISSING = object() # Held while a child is resolved for the first time, so that concurrent @@ -33,6 +31,28 @@ _RESOLUTION_LOCK = threading.RLock() +class _LazyChildren: + """What a node at instant knows about its children. A shallow copy of the + node shares it, as copies of a node that built every child up front shared + their ``_children`` mapping.""" + + __slots__ = ("instant", "source", "resolved", "absent", "added", "children") + + def __init__(self, instant: str, source: dict): + self.instant = instant + # Each child's parameter node as at construction, in order; None once + # every child is resolved. + self.source = source + # Children resolved from ``source`` so far. + self.resolved = {} + # Children of ``source`` with no value at the instant. + self.absent = set() + # Children given with ``add_child``, in call order. + self.added = {} + # Every child in order, once all of ``source`` is resolved. + self.children = None + + class ParameterNodeAtInstant: """ Parameter node of the legislation, at a given instant. @@ -41,9 +61,11 @@ class ParameterNodeAtInstant: descendant up front made each instant a model is asked about cost a copy of the whole parameter tree (16 to 33 MB and several seconds per instant for policyengine-us), which every tax-benefit system then kept for its - lifetime. A child that has not been read yet takes the value its parameter - has when it is first read; a child read once, including one found to have - no value at this instant, keeps what that read found. + lifetime. Which children a node has, and in what order, is fixed when it is + built. A child's value is taken when the child is first read, so a node + held across an in-place change to its parameters reads the changed value + for a child it has not read yet; once read (including found to have no + value), a child keeps what that read found. """ def __init__(self, name: str, node: "ParameterNode", instant_str: str): @@ -56,51 +78,75 @@ def __init__(self, name: str, node: "ParameterNode", instant_str: str): # The "technical" attributes are hidden, so that the node children can be easily browsed with auto-completion without pollution self._name = name self._instant_str = instant_str - state = self.__dict__ - # The parameter node children still resolve from; None once every - # child is resolved. - state[_NODE] = node - state[_INSTANT] = instant_str - # Children resolved so far. Only ever grows, and a shallow copy shares - # it, so every copy reads the same child objects. - state[_RESOLVED] = {} - # Children of the node that have no value at this instant. - state[_ABSENT] = set() - # A child whose name the node itself uses (``add_child``, ``_name``, - # ...) shadows that name, as when every child was set here. - for child_name in _SHADOWING_NAMES: - if child_name in node.children: - _resolve(self, child_name) + source = dict(node.children) + if _STATE in source: + raise ValueError(f"{_STATE!r} cannot name a parameter.") + state = _LazyChildren(instant_str, source) + self.__dict__[_STATE] = state + cls = type(self) + if cls.add_child is not ParameterNodeAtInstant.add_child: + # A subclass's add_child sees every child, as before. + for child_name, child in source.items(): + child_at_instant = child._get_at_instant(instant_str) + if child_at_instant is not None: + self.add_child(child_name, child_at_instant) + with _RESOLUTION_LOCK: + state.source = None + state.children = dict(state.added) + return + for child_name in source: + # A child whose name the node itself answers to shadows it, as + # when every child was set here. + if child_name in _OWN_NAMES or hasattr(cls, child_name): + child_at_instant = _resolve(state, child_name) + if child_at_instant is not _MISSING: + setattr(self, child_name, child_at_instant) def add_child(self, child_name: str, child_at_instant: "ParameterNodeAtInstant"): + if child_name == _STATE: + raise ValueError(f"{_STATE!r} cannot name a parameter.") + state = self.__dict__[_STATE] with _RESOLUTION_LOCK: - _store(self, child_name, child_at_instant) + state.added[child_name] = child_at_instant + if state.children is not None: + state.children[child_name] = child_at_instant + setattr(self, child_name, child_at_instant) def __getattr__(self, key: str): # Reached only for a name that is not an instance attribute: a child - # not read yet, ``_children`` before every child is resolved, or a name - # this node does not have. - if key.startswith("__") and key.endswith("__"): - # Protocol lookups (copy, pickle, numpy) are never parameters. - raise AttributeError(key) - if _RESOLVED not in self.__dict__ or key == "_name": + # not read yet by this instance, ``_children`` before it is set, or a + # name this node does not have. + state = self.__dict__.get(_STATE) + if state is None or key == "_name": # An instance still being rebuilt by copy or pickle. raise AttributeError(key) if key == "_children": - return _materialise(self) - child_at_instant = _resolve(self, key) + return _materialise(self, state) + child_at_instant = _resolve(state, key) if child_at_instant is not _MISSING: + # Attributes hold what the parameters give; ``add_child`` on + # another copy changes only that copy's attribute, as before. + self.__dict__.setdefault(key, child_at_instant) return child_at_instant param_name = helpers._compose_name(self._name, item_name=key) raise ParameterNotFoundError(param_name, self._instant_str) def __dir__(self): - self._children - return [ - name - for name in super().__dir__() - if not name.startswith("_ParameterNodeAtInstant__") - ] + children = self._children + names = set(super().__dir__()) | set(map(str, children)) + names.discard(_STATE) + return list(names) + + def __setstate__(self, state: dict): + self.__dict__.update(state) + if _STATE not in state: + # Pickled when every child was built up front, all in ``_children``. + children = state.get("_children") + if children is None: + children = self.__dict__["_children"] = {} + lazy = _LazyChildren(state.get("_instant_str"), None) + lazy.resolved = lazy.children = children + self.__dict__[_STATE] = lazy def __getitem__( self, key: str @@ -122,11 +168,14 @@ def __getitem__( ) self._vectorial_node = vectorial return vectorial[key] - child_at_instant = self.__dict__[_RESOLVED].get(key, _MISSING) + state = self.__dict__[_STATE] + child_at_instant = state.added.get(key, _MISSING) if child_at_instant is _MISSING: - child_at_instant = _resolve(self, key) + child_at_instant = state.resolved.get(key, _MISSING) if child_at_instant is _MISSING: - raise KeyError(key) + child_at_instant = _resolve(state, key) + if child_at_instant is _MISSING: + raise KeyError(key) return child_at_instant def __iter__(self) -> Iterable: @@ -142,77 +191,69 @@ def __repr__(self) -> str: return result -def _store(node_at_instant, child_name, child_at_instant): - """Record a child. The caller holds ``_RESOLUTION_LOCK``.""" - state = node_at_instant.__dict__ - state[_RESOLVED][child_name] = child_at_instant - state[_ABSENT].discard(child_name) - children = state.get(_CHILDREN) - if children is not None: - children[child_name] = child_at_instant - setattr(node_at_instant, child_name, child_at_instant) - - -def _resolve(node_at_instant, child_name): - """Return one child, resolving it if no read has yet; ``_MISSING`` if the - node has no such child or it has no value at this instant.""" - state = node_at_instant.__dict__ +def _resolve(state: _LazyChildren, child_name): + """A child as the parameters give it, resolving it if no read has yet; + ``_MISSING`` if the node has no such child or it has no value at the + instant.""" with _RESOLUTION_LOCK: - resolved = state[_RESOLVED] - if child_name in resolved: - return resolved[child_name] - if child_name in state[_ABSENT]: - return _MISSING - node = state[_NODE] - child = None if node is None else node.children.get(child_name) + child_at_instant = state.resolved.get(child_name, _MISSING) + if child_at_instant is not _MISSING or child_name in state.absent: + return child_at_instant + source = state.source + child = None if source is None else source.get(child_name) if child is None: return _MISSING - child_at_instant = child._get_at_instant(state[_INSTANT]) - if child_at_instant is None: - state[_ABSENT].add(child_name) + child_at_instant = child._get_at_instant(state.instant) + # A read made while resolving (the same thread re-entering) may have + # recorded this child already; the first recorded result stands. + if child_name in state.absent: + return _MISSING + recorded = state.resolved.get(child_name, _MISSING) + if recorded is not _MISSING: + child_at_instant = recorded + elif child_at_instant is None: + state.absent.add(child_name) else: - _store(node_at_instant, child_name, child_at_instant) - if len(resolved) + len(state[_ABSENT]) >= len(node.children): - # Every child has been read: drop the parameter node. - _materialise(node_at_instant) + state.resolved[child_name] = child_at_instant + if state.source is not None and len(state.resolved) + len(state.absent) == len( + state.source + ): + _complete(state) return _MISSING if child_at_instant is None else child_at_instant -def _materialise(node_at_instant): - """Resolve every child; return them in the parameter node's order, with - any added by ``add_child`` after them.""" - state = node_at_instant.__dict__ - with _RESOLUTION_LOCK: - children = state.get(_CHILDREN) - if children is not None: - return children - node = state[_NODE] - children = {} - for child_name in node.children if node is not None else (): - child_at_instant = _resolve(node_at_instant, child_name) - if child_at_instant is not _MISSING: - children[child_name] = child_at_instant - if state.get(_CHILDREN) is not None: - # Resolving the last child completed the node already. - return state[_CHILDREN] - for child_name, child_at_instant in state[_RESOLVED].items(): - children.setdefault(child_name, child_at_instant) - state[_CHILDREN] = children - if "_children" not in state: - # Unless a child named ``_children`` took the name. - state["_children"] = children - # Every child is resolved: the parameter node is no longer needed. - state[_NODE] = None - return children - - -# Names a ParameterNodeAtInstant itself answers to, which a child of the same -# name shadows. -_SHADOWING_NAMES = frozenset( - {"_name", "_instant_str", "_children", "_vectorial_node"} - | { - name - for name in dir(ParameterNodeAtInstant) - if not name.startswith("__") and not name.startswith("_ParameterNodeAtInstant") +def _complete(state: _LazyChildren): + """Order every child once all of ``source`` is resolved; let go of the + parameter nodes. The caller holds ``_RESOLUTION_LOCK``.""" + resolved, added = state.resolved, state.added + children = { + name: added.get(name, resolved[name]) + for name in state.source + if name in resolved } -) + for name, child_at_instant in added.items(): + # Added children the parameters do not have (or have no value for) + # follow, in the order they were added. + children.setdefault(name, child_at_instant) + if not added: + state.resolved = children + state.children = children + state.source = None + state.absent = set() + + +def _materialise(node_at_instant: ParameterNodeAtInstant, state: _LazyChildren): + """Resolve every child and set ``_children``.""" + with _RESOLUTION_LOCK: + if state.children is None: + for child_name in list(state.source): + if state.children is not None: + break + _resolve(state, child_name) + if state.children is None: + _complete(state) + children = state.children + if "_children" not in node_at_instant.__dict__: + # Unless a child named ``_children`` took the name. + node_at_instant.__dict__["_children"] = children + return children diff --git a/tests/core/parameters/test_parameter_node_at_instant_lazy.py b/tests/core/parameters/test_parameter_node_at_instant_lazy.py index ff2604e02..6ef76d999 100644 --- a/tests/core/parameters/test_parameter_node_at_instant_lazy.py +++ b/tests/core/parameters/test_parameter_node_at_instant_lazy.py @@ -19,7 +19,7 @@ from policyengine_core.country_template import CountryTaxBenefitSystem from policyengine_core.errors import ParameterNotFoundError from policyengine_core.parameters import ParameterNode, ParameterNodeAtInstant -from policyengine_core.parameters.parameter_node_at_instant import _RESOLVED +from policyengine_core.parameters.parameter_node_at_instant import _STATE from policyengine_core.tracers import FullTracer from tests.fixtures.parameter_nodes_at_instant import ( INSTANTS, @@ -49,7 +49,7 @@ def test_country_template_matches_the_up_front_build(instant): def resolved(node_at_instant): - return list(node_at_instant.__dict__[_RESOLVED]) + return list(node_at_instant.__dict__[_STATE].resolved) def test_nothing_is_resolved_until_it_is_read(): @@ -363,3 +363,150 @@ def test_a_node_lets_go_of_its_parameter_node_once_every_child_is_read(): gc.collect() assert tree_ref() is None assert group.x == 1 and list(group) == ["x"] + + +def plain_tree(**values): + return ParameterNode( + "", data={name: {"values": {"2010-01-01": v}} for name, v in values.items()} + ) + + +def test_the_reserved_state_name_is_rejected(): + with pytest.raises(ValueError): + plain_tree(_ParameterNodeAtInstant__state=1).get_at_instant("2017-01-01") + at_instant = plain_tree(x=1).get_at_instant("2017-01-01") + with pytest.raises(ValueError): + at_instant.add_child("_ParameterNodeAtInstant__state", 2) + + +@pytest.mark.parametrize( + "child_name", + [ + "_ParameterNodeAtInstant__resolved", + "_ParameterNodeAtInstant__node", + "_ParameterNodeAtInstant__business", + "__custom__", + "__iter__", + "__dir__", + ], +) +def test_other_unusual_names_are_plain_children(child_name): + at_instant = plain_tree(**{child_name: 42, "ok": 7}).get_at_instant("2017-01-01") + assert getattr(at_instant, child_name) == 42 + assert at_instant[child_name] == 42 + assert at_instant.ok == 7 and at_instant["ok"] == 7 + assert child_name in dir(at_instant) + + +def test_many_added_children_do_not_break_the_first_read(): + names = {f"p{i}": i for i in range(2000)} + at_instant = plain_tree(**names).get_at_instant("2017-01-01") + for i in range(2000): + at_instant.add_child(f"extra{i}", i) + assert at_instant.p0 == 0 + assert len(at_instant._children) == 4000 + assert list(at_instant._children)[:2] == ["p0", "p1"] + + +def test_add_child_on_a_copy_changes_only_shared_items_as_before(): + tree = ParameterNode( + "", + data={ + "x": {"values": {"2010-01-01": 10}}, + "late": {"values": {"2030-01-01": 5}}, + }, + ) + held = tree.get_at_instant("2017-01-01") + with pytest.raises(ParameterNotFoundError): + held.late + duplicate = copy.copy(held) + duplicate.add_child("x", 99) + duplicate.add_child("late", 88) + duplicate.add_child("extra", 77) + # Items are one mapping shared by copies; attributes belong to each copy. + assert held["x"] == 99 and held.x == 10 + assert held["late"] == 88 + with pytest.raises(ParameterNotFoundError): + held.late + assert held["extra"] == 77 + with pytest.raises(ParameterNotFoundError): + held.extra + assert duplicate.x == 99 and duplicate.late == 88 and duplicate.extra == 77 + assert list(held._children) == ["x", "late", "extra"] + + +def test_a_copy_installs_what_it_reads_as_attributes(): + held = build_tree().get_at_instant("2017-03-01") + duplicate = copy.copy(held) + scale = held.scale + assert duplicate.scale is scale + assert duplicate.__dict__["scale"] is scale + + +class _MethodCollision(ParameterNodeAtInstant): + def policy(self): + return "method" + + +class _DoublingAddChild(ParameterNodeAtInstant): + def add_child(self, child_name, child_at_instant): + super().add_child(child_name, child_at_instant * 2) + + +def test_subclasses_keep_their_collisions_and_add_child_hooks(): + tree = plain_tree(policy=42, x=5) + collided = _MethodCollision("", tree, "2017-01-01") + assert collided.policy == 42 and collided["policy"] == 42 + doubled = _DoublingAddChild("", tree, "2017-01-01") + assert doubled.x == 10 and doubled["x"] == 10 + assert dict(doubled._children) == {"policy": 84, "x": 10} + + +class ReenteringChild: + """A child whose first resolution reads the same child again.""" + + def __init__(self): + self.node_at_instant = None + self.calls = 0 + + def _get_at_instant(self, instant): + self.calls += 1 + result = object() + if self.calls == 1: + self.inner = self.node_at_instant["again"] + return result + + +def test_a_reentrant_read_keeps_the_first_recorded_child(): + tree = plain_tree(other=1) + child = ReenteringChild() + tree.children = {"again": child, **tree.children} + at_instant = tree.get_at_instant("2017-01-01") + child.node_at_instant = at_instant + outer = at_instant["again"] + assert outer is child.inner + assert at_instant["again"] is outer and at_instant.again is outer + + +def test_children_and_their_order_are_fixed_when_the_node_is_built(): + tree = plain_tree(x=10, y=20, z=30) + held = tree.get_at_instant("2017-01-01") + assert held.x == 10 + del tree.children["x"] + tree.children["w"] = ParameterNode("", data={"v": {"values": {"2010-01-01": 1}}}) + assert list(held) == ["x", "y", "z"] + assert held.x == 10 + assert "w" not in held._children + + +def test_a_node_pickled_when_every_child_was_built_up_front_loads(): + """State as the up-front build pickled it: everything in ``_children``.""" + restored = ParameterNodeAtInstant.__new__(ParameterNodeAtInstant) + restored.__setstate__( + {"_name": "", "_instant_str": "2017-01-01", "_children": {"x": 10}, "x": 10} + ) + assert restored.x == 10 and restored["x"] == 10 and list(restored) == ["x"] + restored.add_child("y", 2) + assert restored["y"] == 2 and list(restored) == ["x", "y"] + with pytest.raises(KeyError): + restored["nope"]