Read only periods the current branch can see when uprating or carrying over

[MaxGhenis]

Follow-up to #551, whose review flagged this as an existing issue that #551 left unchanged.

Problem

Holder.get_known_periods() lists the period of every stored key with the branch name stripped. Holder.get_array(period, branch_name) reads only that branch, its parent_branch ancestors and default. Simulation._calculate took the latest known period from the unscoped list and read it back under self.branch_name, so a period stored only under an unreadable branch came back as None:

• Uprating computed None * factor → TypeError: unsupported operand type(s) for *: 'NoneType' and 'float'.
• Auto-carry-over passed None to _cast_formula_result, which cached NaN with no error.

The same gap showed up in two more places, all confirmed by running them on master:

• OnDiskStorage parsed f"{branch}_{period}" keys with split("_")[1]. policyengine-us branch names such as no_salt, mtr_for_adult_1 and pre_tcja_ctc contain _:
  • no_salt → period "salt" → ValueError.
  • A name like y_2019 → period 2019, a period that was never stored. Uprating then picked it and hit the TypeError above.
  • get_known_branch_periods raised "too many values to unpack".
  • delete(None, "pre_tcja") prefix-matched and also wiped pre_tcja_ctc.
• dump_simulation read every listed period under default. Dumping a branch saved None for its branch-only periods (and default's value where the branch overrode it). restore_simulation then failed with "Object arrays cannot be loaded when allow_pickle=False".

How a holder gets an unreadable period

• On disk: reachable through public API only. Set simulation.memory_config, then create a holder after it (for example with a variable added to the system after the simulation is built). A branch whose name contains _ then stores its input on disk, and branch.calculate for a later year raises. Simulation.__init__ creates a holder for every existing variable before memory_config can be set, so on-disk storage needs that setup. policyengine-us doesn't set memory_config (git grep).
• In memory: needs a holder-level call that names another branch, such as Holder.set_input(period, array, branch_name) or put_in_cache(..., branch_name). Through Simulation methods I found no construction that does it: sibling branches, nested branches, and set_input on a branch after _user_input_contexts exists. The review suggested Holder._set writing under _user_input_contexts[-1] as a route. That list is shared by reference across Simulation.clone, but a branch's set_input wrote only that branch's own holder. An independent search for other constructions is running; I'll update this section with its result.

Fix

• One definition of what a branch reads. Holder._readable_branch_names(branch_name) returns the branch, then (unless it is default) each parent_branch ancestor, then default. get_array uses it, and the lookup order is unchanged; the branch's own value is still checked first, before the list is built.
• Holder.get_known_periods(branch_name=None). With a branch name, it lists only the periods that branch can read. Without one, it behaves as before.
• _calculate calls holder.get_known_periods(self.branch_name). Uprating and carry-over now use the latest period this branch can read; with none, they fall through to the default value.
• OnDiskStorage splits keys on the last _ (a period's string form never contains _; see Period.__str__). delete(None, branch) compares the parsed branch name exactly.
• dump_simulation saves, once per period, the value the dumped simulation's branch reads.

Invariants

For every holder h, branch b and stored state:

1. Readability agreement. set(h.get_known_periods(b)) == {p : h.get_array(p, b) is not None}, with no duplicates. Checked exhaustively on every non-empty subset of stored branches {default, a, a_b, c}, in memory and on disk. The readers are the root, a, a_b, the sibling c, and a branch literally named default nested under a.
2. Isolation. Values stored only under branches a simulation cannot read never change what it calculates. This is a differential check against a simulation that never had them: 448 cases covering uprating and carry-over, 4 readable-year sets × 14 unreadable-year sets × 4 requested years.
3. Finite results. Neither path returns NaN or raises because of an unreadable period.
4. Disk key round trip. For branch names with and without _ and every period form (year, month, day, multi-unit, anchored year:2025-03, ETERNITY), get_known_branch_periods returns exactly what was put, get reads it back, and restore() rebuilds the same index.
5. Unchanged elsewhere. With no unreadable periods, results match master. The existing suite passes unchanged.

Tests

tests/core/test_known_periods_branch_visibility.py (new, 532 cases). On master, 280 of them fail:

Cause on master	Failing cases
Uprating TypeError (None * float)	92
Carry-over NaN	54
Wrong value (isolation check and a later unreadable period)	65
On-disk ValueError (unpack / period parse)	44
Dump object-array load failure	1
Delete prefix collision	1
Call to the new get_known_periods(branch_name) argument	23

The 252 that pass on master are the cases with nothing unreadable stored, and nested branches reading ancestor periods. They confirm invariant 5.

Run locally on this branch:

• uv run pytest tests: 1477 passed, 4 skipped, 1 xfailed.
• Country-template YAML (policyengine-core test policyengine_core/country_template/tests -c policyengine_core.country_template): 39 passed.
• uvx ruff format --check .: clean.
• Not run: make documentation. The Holder page autodocs its members, so the updated get_array and get_known_periods docstrings are the documentation change.

Grid tests rather than Hypothesis, as in #551: Hypothesis isn't a dependency, and these domains are small enough to enumerate.

axiom: n/a: engine fix to branch storage lookups in policyengine-core; no policy rule changes

🤖 Generated with Claude Code