From bbb4f644019eef05611ff17dc9f1f4ef3ba545af Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Sun, 27 Sep 2026 15:50:27 -0400 Subject: [PATCH 01/15] Read only periods the current branch can see when uprating or carrying over Holder.get_known_periods() lists the periods of every stored key, with the branch name stripped, but Holder.get_array() reads only the requested branch, its parent_branch ancestors and "default". Simulation._calculate took the latest known period from the unscoped list, so a period stored only under an unrelated branch read back as None: uprating raised TypeError and auto-carry-over cached NaN. - Holder._readable_branch_names() is the one definition of what a branch can read; get_array() and the new get_known_periods(branch_name) both use it. - _calculate uses get_known_periods(self.branch_name). - OnDiskStorage splits "_" keys on the last "_": branch names like "no_salt" raised ValueError, "y_2019" listed the wrong period, and delete(None, "pre_tcja") also wiped "pre_tcja_ctc". - dump_simulation saves the values the dumped branch reads instead of reading every period under "default" and saving None. Co-Authored-By: Claude Opus 5.5 --- .../fix-uprating-branch-visibility.fixed.md | 1 + .../data_storage/on_disk_storage.py | 22 +- policyengine_core/holders/holder.py | 82 ++-- policyengine_core/simulations/simulation.py | 7 +- policyengine_core/tools/simulation_dumper.py | 12 +- .../test_known_periods_branch_visibility.py | 422 ++++++++++++++++++ 6 files changed, 504 insertions(+), 42 deletions(-) create mode 100644 changelog.d/fix-uprating-branch-visibility.fixed.md create mode 100644 tests/core/test_known_periods_branch_visibility.py diff --git a/changelog.d/fix-uprating-branch-visibility.fixed.md b/changelog.d/fix-uprating-branch-visibility.fixed.md new file mode 100644 index 000000000..36eb6f762 --- /dev/null +++ b/changelog.d/fix-uprating-branch-visibility.fixed.md @@ -0,0 +1 @@ +Uprating and auto-carry-over now use only periods the current branch can read (its own, its parent branches' and `default`), so a period stored only under an unrelated branch no longer raises `TypeError` or caches `NaN`; on-disk storage parses branch names that contain `_`, and `dump_simulation` saves the values the dumped branch reads. diff --git a/policyengine_core/data_storage/on_disk_storage.py b/policyengine_core/data_storage/on_disk_storage.py index 3563c805d..3de90e957 100644 --- a/policyengine_core/data_storage/on_disk_storage.py +++ b/policyengine_core/data_storage/on_disk_storage.py @@ -9,6 +9,17 @@ from policyengine_core.periods import Period +def _split_key(key: str) -> tuple: + """Split a ``f"{branch_name}_{period}"`` file key into its two parts. + + Branch names often contain ``_`` (policyengine-us uses ``no_salt`` and + ``mtr_for_adult_1``) but a period's string form never does, so the key + splits on its last ``_``. + """ + branch_name, period = key.rsplit("_", 1) + return branch_name, period + + class OnDiskStorage: """ Low-level class responsible for storing and retrieving calculated vectors on disk @@ -82,12 +93,13 @@ def delete(self, period: Period = None, branch_name: str = "default") -> None: if period is None: # Only wipe files belonging to the requested branch (previously # this wiped every branch regardless of ``branch_name`` — same - # class of bug as C2 in InMemoryStorage). - branch_prefix = f"{branch_name}_" + # class of bug as C2 in InMemoryStorage). Compare the parsed + # branch name, not a prefix: deleting ``pre_tcja`` must not + # also wipe ``pre_tcja_ctc``. self._files = { period_item: value for period_item, value in self._files.items() - if not period_item.startswith(branch_prefix) + if _split_key(period_item)[0] != branch_name } return @@ -103,12 +115,12 @@ def delete(self, period: Period = None, branch_name: str = "default") -> None: } def get_known_periods(self) -> list: - return list([periods.period(x.split("_")[1]) for x in self._files.keys()]) + return [period for _, period in self.get_known_branch_periods()] def get_known_branch_periods(self) -> list: return [ (branch_name, periods.period(period)) - for branch_name, period in map(lambda x: x.split("_"), self._files.keys()) + for branch_name, period in map(_split_key, self._files.keys()) ] def restore(self) -> None: diff --git a/policyengine_core/holders/holder.py b/policyengine_core/holders/holder.py index 52ae81295..96b4675b6 100644 --- a/policyengine_core/holders/holder.py +++ b/policyengine_core/holders/holder.py @@ -106,45 +106,53 @@ def _get_array_from_storage( value = self._disk_storage.get(period, branch_name) return value + def _readable_branch_names(self, branch_name: str = "default") -> List[str]: + """ + Branches whose stored values ``get_array(period, branch_name)`` can + return, in lookup order. + + That is ``branch_name`` itself, then (unless it is ``default``) each + ``simulation.parent_branch`` ancestor, then ``default``. Nested + branches inherit values from their parent (e.g. a ``no_salt`` branch + cloned from an ``itemizing`` branch still sees ``tax_unit_itemizes`` + set on the ``itemizing`` branch). Previously the fallback returned + the first branch in dict-insertion order (bug C1) — silently swapping + values between unrelated sibling branches (reform vs baseline) and + producing wrong reform deltas. The post-C1 behavior only fell back to + ``default``, which broke country-package nested-branch patterns that + relied on the ancestor's input being visible. + """ + if branch_name == "default": + return ["default"] + branch_names = [branch_name] + simulation = getattr(self, "simulation", None) + parent = getattr(simulation, "parent_branch", None) if simulation else None + while parent is not None: + branch_names.append(parent.branch_name) + parent = getattr(parent, "parent_branch", None) + branch_names.append("default") + return list(dict.fromkeys(branch_names)) + def get_array(self, period: Period, branch_name: str = "default") -> ArrayLike: """ Get the value of the variable for the given period. - If the value is not known, return ``None``. + Values stored under ``branch_name``, its ``parent_branch`` ancestors + and ``default`` are visible, in that order. If the value is not + known on any of them, return ``None``. """ if self.variable.is_neutralized: return self.default_array() + # The branch's own value is the common case: look it up before + # walking the ancestors. value = self._get_array_from_storage(period, branch_name) if value is not None: return value - if value is None and branch_name != "default": - # Walk up ``simulation.parent_branch`` so nested branches inherit - # values from their parent (e.g. a ``no_salt`` branch cloned - # from an ``itemizing`` branch still sees ``tax_unit_itemizes`` - # set on the ``itemizing`` branch). Fall back to ``default`` - # only if no ancestor branch has a value. Previously the - # fallback returned the first branch in dict-insertion order - # (bug C1) — silently swapping values between unrelated - # sibling branches (reform vs baseline) and producing wrong - # reform deltas. The post-C1 behavior only fell back to - # ``default``, which broke country-package nested-branch - # patterns that relied on the ancestor's input being visible. - parent = ( - getattr(self.simulation, "parent_branch", None) - if self.simulation - else None - ) - while parent is not None: - ancestor_value = self._get_array_from_storage( - period, - parent.branch_name, - ) - if ancestor_value is not None: - return ancestor_value - parent = getattr(parent, "parent_branch", None) - default_value = self._get_array_from_storage(period, "default") - if default_value is not None: - return default_value + for readable_branch_name in self._readable_branch_names(branch_name)[1:]: + value = self._get_array_from_storage(period, readable_branch_name) + if value is not None: + return value + return None def get_memory_usage(self) -> dict: """ @@ -189,11 +197,23 @@ def get_memory_usage(self) -> dict: return usage - def get_known_periods(self) -> List[Period]: + def get_known_periods(self, branch_name: str = None) -> List[Period]: """ Get the list of periods the variable value is known for. - """ + With ``branch_name``, list only the periods ``get_array(period, + branch_name)`` can read: those stored under that branch, its + ``parent_branch`` ancestors or ``default``. Without it, list the + periods stored under every branch in this holder, some of which that + branch may not be able to read. + """ + if branch_name is not None: + readable_branch_names = set(self._readable_branch_names(branch_name)) + return [ + period + for stored_branch_name, period in self.get_known_branch_periods() + if stored_branch_name in readable_branch_names + ] return list(self._memory_storage.get_known_periods()) + list( (self._disk_storage.get_known_periods() if self._disk_storage else []) ) diff --git a/policyengine_core/simulations/simulation.py b/policyengine_core/simulations/simulation.py index 566407a59..bc3a13483 100644 --- a/policyengine_core/simulations/simulation.py +++ b/policyengine_core/simulations/simulation.py @@ -864,8 +864,11 @@ def _calculate(self, variable_name: str, period: Period = None) -> ArrayLike: # If no result, use the default value and cache it if array is None: - # Check if the variable has a previously defined value - known_periods = holder.get_known_periods() + # Check if the variable has a previously defined value. + # Only periods this branch can read count: a period stored + # only under an unrelated branch would read back as ``None`` + # and reach the arithmetic below. + known_periods = holder.get_known_periods(self.branch_name) earlier_known_periods = [ known_period for known_period in known_periods diff --git a/policyengine_core/tools/simulation_dumper.py b/policyengine_core/tools/simulation_dumper.py index c3db0c4fa..b22da1a64 100644 --- a/policyengine_core/tools/simulation_dumper.py +++ b/policyengine_core/tools/simulation_dumper.py @@ -32,7 +32,7 @@ def dump_simulation(simulation, directory): # Dump variable values for holder in entity._holders.values(): - _dump_holder(holder, directory) + _dump_holder(holder, directory, simulation.branch_name) def restore_simulation(directory, tax_benefit_system, **kwargs): @@ -64,10 +64,14 @@ def restore_simulation(directory, tax_benefit_system, **kwargs): return simulation -def _dump_holder(holder, directory): +def _dump_holder(holder, directory, branch_name="default"): + # Dump the values the simulation's branch reads. A holder can also store + # periods under branches this one cannot see; reading those back under + # this branch gives ``None``, which would be saved as an object array + # that ``restore_simulation`` cannot load. disk_storage = holder.create_disk_storage(directory, preserve=True) - for period in holder.get_known_periods(): - value = holder.get_array(period) + for period in dict.fromkeys(holder.get_known_periods(branch_name)): + value = holder.get_array(period, branch_name) disk_storage.put(value, period) diff --git a/tests/core/test_known_periods_branch_visibility.py b/tests/core/test_known_periods_branch_visibility.py new file mode 100644 index 000000000..43ae823dc --- /dev/null +++ b/tests/core/test_known_periods_branch_visibility.py @@ -0,0 +1,422 @@ +"""Known periods a branch cannot read. + +Holder storage keys embed the branch name (``"no_salt:2018"`` in memory, +``"no_salt_2018"`` on disk), and a holder can hold keys for branches other +than the simulation's own. ``Holder.get_known_periods()`` lists the periods of +every key, while ``Holder.get_array(period, branch_name)`` reads only that +branch, its ``parent_branch`` ancestors and ``default``. Three defects came +from that gap: + +* ``Simulation._calculate`` took the latest known earlier period from the + unscoped list. If that period was stored only under a branch the simulation + cannot read, ``get_array`` returned ``None``: uprating raised ``TypeError`` + (``None * factor``) and auto-carry-over cached ``NaN``. +* ``OnDiskStorage`` parsed ``f"{branch}_{period}"`` keys with + ``split("_")[1]``. A branch name containing ``_`` raised ``ValueError`` + (``no_salt`` -> period ``"salt"``) or listed the wrong period + (``y_2019`` -> ``2019``). ``delete(None, "pre_tcja")`` also wiped + ``pre_tcja_ctc``, whose key shares the prefix. +* ``dump_simulation`` read every listed period under ``default``, saving + ``None`` for a period stored only on the dumped branch, which + ``restore_simulation`` could not load. +""" + +import itertools +import math +import os +import tempfile +import warnings + +import numpy as np +import pytest + +from policyengine_core import periods +from policyengine_core.country_template import CountryTaxBenefitSystem +from policyengine_core.country_template.entities import Person +from policyengine_core.data_storage import OnDiskStorage +from policyengine_core.experimental import MemoryConfig +from policyengine_core.model_api import YEAR, Variable +from policyengine_core.parameters import ParameterNode +from policyengine_core.simulations import SimulationBuilder +from policyengine_core.tools import simulation_dumper + +INDEX_START = 2015 +# The index grows 10% a year, so ratios between years are powers of 1.1. +INDEX = { + f"{year}-01-01": 100 * 1.1 ** (year - INDEX_START) for year in range(2015, 2023) +} + + +def growth(from_year: int, to_year: int) -> float: + return 1.1 ** (to_year - from_year) + + +def build_system(auto_carry_over: bool = False) -> CountryTaxBenefitSystem: + system = CountryTaxBenefitSystem() + system.auto_carry_over_input_variables = auto_carry_over + system.parameters.add_child( + "test_uprating", + ParameterNode("test_uprating", data={"index": {"values": INDEX}}), + ) + + class uprated_income(Variable): + value_type = float + entity = Person + definition_period = YEAR + label = "Uprated yearly income" + uprating = "test_uprating.index" + + class carried_income(Variable): + value_type = float + entity = Person + definition_period = YEAR + label = "Yearly income carried over without uprating" + + system.add_variable(uprated_income) + system.add_variable(carried_income) + return system + + +@pytest.fixture(scope="module") +def system(): + return build_system() + + +@pytest.fixture(scope="module") +def carry_over_system(): + return build_system(auto_carry_over=True) + + +def new_simulation(system): + return SimulationBuilder().build_default_simulation(system, count=1) + + +def store(simulation, variable, year, value, branch_name): + """Store ``value`` for ``year`` under ``branch_name`` in the holder of + ``simulation``, the way ``Holder.set_input`` does when a caller names a + branch other than the simulation's own.""" + simulation.get_holder(variable).set_input( + periods.period(year), np.array([value]), branch_name + ) + + +def only(array) -> float: + (value,) = array + return float(value) + + +# ----- Simulation._calculate -------------------------------------------------- + + +def test_uprating_skips_period_stored_only_on_unrelated_branch(system): + simulation = new_simulation(system) + store(simulation, "uprated_income", 2016, 1_000.0, "default") + store(simulation, "uprated_income", 2018, 5_000.0, "other") + + result = simulation.calculate("uprated_income", 2020) + + # Uprated from 2016, the latest year the default branch can read. + assert only(result) == pytest.approx(1_000 * growth(2016, 2020)) + + +def test_uprating_with_no_readable_earlier_period_uses_default(system): + simulation = new_simulation(system) + store(simulation, "uprated_income", 2018, 5_000.0, "other") + + assert only(simulation.calculate("uprated_income", 2020)) == 0 + + +def test_uprating_in_branch_skips_sibling_branch_period(system): + simulation = new_simulation(system) + simulation.set_input("uprated_income", 2016, [1_000.0]) + branch = simulation.get_branch("reform") + store(branch, "uprated_income", 2018, 5_000.0, "baseline") + + result = branch.calculate("uprated_income", 2020) + + assert only(result) == pytest.approx(1_000 * growth(2016, 2020)) + + +def test_uprating_in_nested_branch_reads_ancestor_periods(system): + simulation = new_simulation(system) + simulation.set_input("uprated_income", 2016, [1_000.0]) + itemizing = simulation.get_branch("itemizing") + itemizing.set_input("uprated_income", 2018, [5_000.0]) + no_salt = itemizing.get_branch("no_salt") + + result = no_salt.calculate("uprated_income", 2020) + + assert only(result) == pytest.approx(5_000 * growth(2018, 2020)) + + +def test_carry_over_skips_period_stored_only_on_unrelated_branch( + carry_over_system, +): + simulation = new_simulation(carry_over_system) + store(simulation, "carried_income", 2016, 1_000.0, "default") + store(simulation, "carried_income", 2018, 5_000.0, "other") + + result = simulation.calculate("carried_income", 2020) + + assert not math.isnan(only(result)) + assert only(result) == 1_000 + + +def test_carry_over_with_no_readable_period_uses_default(carry_over_system): + simulation = new_simulation(carry_over_system) + store(simulation, "carried_income", 2018, 5_000.0, "other") + + result = simulation.calculate("carried_income", 2020) + + assert not math.isnan(only(result)) + assert only(result) == 0 + + +def test_carry_over_ignores_later_period_on_unrelated_branch(carry_over_system): + simulation = new_simulation(carry_over_system) + store(simulation, "carried_income", 2016, 1_000.0, "default") + store(simulation, "carried_income", 2022, 9_000.0, "other") + + assert only(simulation.calculate("carried_income", 2020)) == 1_000 + + +# Invariant: values stored only under branches a simulation cannot read do not +# change what it calculates. Checked exhaustively against a simulation that +# never had those values, for every combination of readable years, unreadable +# years and requested year below. +READABLE_YEAR_SETS = [(), (2016,), (2018,), (2016, 2018)] +UNREADABLE_YEAR_SETS = [ + years + for size in range(1, 4) + for years in itertools.combinations((2015, 2017, 2019, 2021), size) +] +REQUESTED_YEARS = (2017, 2018, 2019, 2020) + + +def _calculate_with(system, variable, readable, unreadable, requested): + simulation = new_simulation(system) + branch = simulation.get_branch("reform") + for year in readable: + store(branch, variable, year, 1_000.0 * (year - 2000), "reform") + for year in unreadable: + store(branch, variable, year, 7_777.0, "baseline") + return only(branch.calculate(variable, requested)) + + +@pytest.mark.parametrize("requested", REQUESTED_YEARS) +@pytest.mark.parametrize("unreadable", UNREADABLE_YEAR_SETS) +@pytest.mark.parametrize("readable", READABLE_YEAR_SETS) +@pytest.mark.parametrize( + "variable, auto_carry_over", + [("uprated_income", False), ("carried_income", True)], +) +def test_unreadable_branch_periods_do_not_change_results( + system, + carry_over_system, + variable, + auto_carry_over, + readable, + unreadable, + requested, +): + tax_benefit_system = carry_over_system if auto_carry_over else system + + with_unreadable = _calculate_with( + tax_benefit_system, variable, readable, unreadable, requested + ) + without_unreadable = _calculate_with( + tax_benefit_system, variable, readable, (), requested + ) + + assert math.isfinite(with_unreadable) + assert with_unreadable == without_unreadable + + +# ----- Holder.get_known_periods(branch_name) agrees with get_array ------------ + +# default -> a -> a_b is one lineage; c is a sibling of a, and a +# branch named "default" nested under a reads only "default" in get_array. +STORED_BRANCHES = ("default", "a", "a_b", "c") + + +def _lineage(system): + simulation = new_simulation(system) + a = simulation.get_branch("a") + return { + "default": simulation, + "a": a, + "a_b": a.get_branch("a_b"), + "c": simulation.get_branch("c"), + "default under a": a.get_branch("default"), + } + + +@pytest.mark.parametrize("on_disk", [False, True]) +@pytest.mark.parametrize( + "stored", + [ + branches + for size in range(1, len(STORED_BRANCHES) + 1) + for branches in itertools.combinations(STORED_BRANCHES, size) + ], +) +def test_known_periods_for_branch_are_exactly_the_readable_ones( + system, tmp_path, stored, on_disk +): + """For every reader, ``get_known_periods(branch_name)`` lists a period if + and only if ``get_array(period, branch_name)`` returns a value.""" + for reader_name, reader in _lineage(system).items(): + holder = reader.get_holder("uprated_income") + if on_disk: + directory = tmp_path / reader_name.replace(" ", "-") + directory.mkdir() + holder._disk_storage = holder.create_disk_storage( + str(directory), preserve=True + ) + storage = holder._disk_storage if on_disk else holder._memory_storage + # One distinct year per stored branch, so each listed period traces + # back to exactly one branch. + for offset, branch_name in enumerate(stored): + storage.put(np.array([1.0]), periods.period(2016 + offset), branch_name) + + stored_periods = {period for _, period in holder.get_known_branch_periods()} + readable = { + period + for period in stored_periods + if holder.get_array(period, reader.branch_name) is not None + } + + listed = holder.get_known_periods(reader.branch_name) + + assert set(listed) == readable, reader_name + assert len(listed) == len(set(listed)), reader_name + + +def test_get_known_periods_without_branch_lists_every_branch(system): + simulation = new_simulation(system) + store(simulation, "uprated_income", 2016, 1.0, "default") + store(simulation, "uprated_income", 2018, 1.0, "other") + + holder = simulation.get_holder("uprated_income") + + assert sorted(holder.get_known_periods()) == [ + periods.period(2016), + periods.period(2018), + ] + assert holder.get_known_periods("default") == [periods.period(2016)] + + +# ----- OnDiskStorage key parsing ---------------------------------------------- + +BRANCH_NAMES = ( + "default", + "no_salt", + "mtr_for_adult_1", + "pre_tcja_ctc", + "y_2019", + "trailing_", +) +PERIODS = ( + periods.period(2025), + periods.period("2025-03"), + periods.period("2025-03-05"), + periods.period("month:2025-01:3"), + periods.period("year:2025-03"), + periods.period("year:2024:2"), +) + + +@pytest.fixture +def disk_storage(tmp_path): + return OnDiskStorage(str(tmp_path), preserve_storage_dir=True) + + +@pytest.mark.parametrize("period", PERIODS, ids=str) +@pytest.mark.parametrize("branch_name", BRANCH_NAMES) +def test_on_disk_keys_round_trip(disk_storage, branch_name, period): + disk_storage.put(np.array([3.0]), period, branch_name) + + assert disk_storage.get_known_branch_periods() == [(branch_name, period)] + assert disk_storage.get_known_periods() == [period] + np.testing.assert_array_equal(disk_storage.get(period, branch_name), [3.0]) + + # The same holds after rebuilding the index from the files on disk. + disk_storage.restore() + assert disk_storage.get_known_branch_periods() == [(branch_name, period)] + + +@pytest.mark.parametrize("branch_name", BRANCH_NAMES) +def test_on_disk_eternal_keys_round_trip(tmp_path, branch_name): + storage = OnDiskStorage(str(tmp_path), is_eternal=True, preserve_storage_dir=True) + storage.put(np.array([3.0]), periods.period(2025), branch_name) + + assert storage.get_known_branch_periods() == [ + (branch_name, periods.period(periods.ETERNITY)) + ] + + +def test_on_disk_delete_branch_leaves_branches_sharing_its_prefix(disk_storage): + period = periods.period(2025) + for branch_name in ("pre_tcja", "pre_tcja_ctc", "pre"): + disk_storage.put(np.array([1.0]), period, branch_name) + + disk_storage.delete(None, "pre_tcja") + + assert sorted(disk_storage.get_known_branch_periods()) == [ + ("pre", period), + ("pre_tcja_ctc", period), + ] + + +@pytest.mark.parametrize("branch_name", ["no_salt", "y_2019"]) +def test_disk_backed_branch_with_underscore_uprates(branch_name): + """Only public API: a variable added after ``memory_config`` is set gets + disk storage, and the branch's input is uprated from its disk key.""" + tax_benefit_system = build_system() + simulation = new_simulation(tax_benefit_system) + with warnings.catch_warnings(): + warnings.simplefilter("ignore") + simulation.memory_config = MemoryConfig(max_memory_occupation=0) + + class disk_income(Variable): + value_type = float + entity = Person + definition_period = YEAR + label = "Uprated yearly income stored on disk" + uprating = "test_uprating.index" + + tax_benefit_system.add_variable(disk_income) + branch = simulation.get_branch(branch_name) + branch.set_input("disk_income", 2018, [5_000.0]) + disk_keys = list(branch.get_holder("disk_income")._disk_storage._files) + assert disk_keys == [f"{branch_name}_2018"] + + result = branch.calculate("disk_income", 2020) + + assert only(result) == pytest.approx(5_000 * growth(2018, 2020)) + + +# ----- dump_simulation -------------------------------------------------------- + + +def test_dump_branch_saves_the_values_the_branch_reads(system): + simulation = new_simulation(system) + simulation.set_input("uprated_income", 2016, [1_000.0]) + simulation.set_input("uprated_income", 2017, [2_000.0]) + branch = simulation.get_branch("reform") + branch.set_input("uprated_income", 2017, [3_000.0]) + branch.set_input("uprated_income", 2018, [5_000.0]) + + directory = os.path.join(tempfile.mkdtemp(), "dump") + simulation_dumper.dump_simulation(branch, directory) + restored = simulation_dumper.restore_simulation(directory, system) + + holder = restored.get_holder("uprated_income") + assert sorted(holder.get_known_periods()) == [ + periods.period(2016), + periods.period(2017), + periods.period(2018), + ] + assert only(holder.get_array(2016)) == 1_000 + assert only(holder.get_array(2017)) == 3_000 + assert only(holder.get_array(2018)) == 5_000 From 5fe8b61864a3a929d3394e1658a83de1f4ccf7be Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Fri, 2 Oct 2026 07:44:44 -0400 Subject: [PATCH 02/15] Skip on-disk round trips of ':' periods on Windows The string form of a multi-unit or rolling-year period contains ':', which a Windows file name cannot. On windows-latest, numpy.save raised OSError for month:2025-01:3 and year:2024:2, and year:2025-03 read back until restore() found no file for it. Disk storage has never kept these periods on Windows; this PR does not change that, so the test skips them there. Co-Authored-By: Claude Opus 5.5 --- tests/core/test_known_periods_branch_visibility.py | 12 +++++++++--- 1 file changed, 9 insertions(+), 3 deletions(-) diff --git a/tests/core/test_known_periods_branch_visibility.py b/tests/core/test_known_periods_branch_visibility.py index 43ae823dc..d1806af9b 100644 --- a/tests/core/test_known_periods_branch_visibility.py +++ b/tests/core/test_known_periods_branch_visibility.py @@ -316,13 +316,19 @@ def test_get_known_periods_without_branch_lists_every_branch(system): "y_2019", "trailing_", ) +# A multi-unit or rolling-year period's string form contains ``:``, which a +# Windows file name cannot, so disk storage cannot keep those periods there. +skip_on_windows = pytest.mark.skipif( + os.name == "nt", reason="Windows file names cannot contain ':'." +) PERIODS = ( periods.period(2025), periods.period("2025-03"), periods.period("2025-03-05"), - periods.period("month:2025-01:3"), - periods.period("year:2025-03"), - periods.period("year:2024:2"), + *( + pytest.param(periods.period(period), id=period, marks=skip_on_windows) + for period in ("month:2025-01:3", "year:2025-03", "year:2024:2") + ), ) From bfa78cba8f229f9029dd17fe6609e8f91644c2a3 Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Fri, 2 Oct 2026 07:33:23 -0400 Subject: [PATCH 03/15] Delete the periods within a deleted period from disk storage OnDiskStorage.delete(period, branch_name) removed only the file keyed exactly f"{branch_name}_{period}". InMemoryStorage.delete and the Holder.delete_arrays docstring remove every period the given period contains, so a disk-backed holder kept values the caller had deleted (deleting "2025" left "2025-01" and "2025-02"). Parse each key with _split_key and delete those whose branch is branch_name and whose period the deleted period contains. Eternal storage still deletes its one ETERNITY key. Fixes #564. Co-Authored-By: Claude Opus 5.5 --- ...fix-disk-delete-contained-periods.fixed.md | 1 + .../data_storage/on_disk_storage.py | 15 +- ...t_disk_storage_delete_contained_periods.py | 318 ++++++++++++++++++ .../core/test_disk_storage_delete_property.py | 121 +++++++ 4 files changed, 454 insertions(+), 1 deletion(-) create mode 100644 changelog.d/fix-disk-delete-contained-periods.fixed.md create mode 100644 tests/core/test_disk_storage_delete_contained_periods.py create mode 100644 tests/core/test_disk_storage_delete_property.py diff --git a/changelog.d/fix-disk-delete-contained-periods.fixed.md b/changelog.d/fix-disk-delete-contained-periods.fixed.md new file mode 100644 index 000000000..9748a6f83 --- /dev/null +++ b/changelog.d/fix-disk-delete-contained-periods.fixed.md @@ -0,0 +1 @@ +On-disk storage now deletes every period the deleted period contains, as in-memory storage does, so deleting `2025` from a disk-backed holder also deletes `2025-01` through `2025-12`. diff --git a/policyengine_core/data_storage/on_disk_storage.py b/policyengine_core/data_storage/on_disk_storage.py index 3de90e957..69dc52abd 100644 --- a/policyengine_core/data_storage/on_disk_storage.py +++ b/policyengine_core/data_storage/on_disk_storage.py @@ -20,6 +20,15 @@ def _split_key(key: str) -> tuple: return branch_name, period +def _is_within(key: str, branch_name: str, period: Period) -> bool: + """Whether file key ``key`` stores a value of ``branch_name`` for a + period that ``period`` contains.""" + key_branch_name, key_period = _split_key(key) + return key_branch_name == branch_name and period.contains( + periods.period(key_period) + ) + + class OnDiskStorage: """ Low-level class responsible for storing and retrieving calculated vectors on disk @@ -107,11 +116,15 @@ def delete(self, period: Period = None, branch_name: str = "default") -> None: period = periods.period(periods.ETERNITY) period = periods.period(period) + # Delete every period of the branch that ``period`` contains, as + # ``InMemoryStorage.delete`` does and ``Holder.delete_arrays`` + # documents: deleting ``2025`` also deletes ``2025-01``. Previously + # only the file keyed exactly ``period`` was deleted. if period is not None: self._files = { period_item: value for period_item, value in self._files.items() - if not period_item == f"{branch_name}_{period}" + if not _is_within(period_item, branch_name, period) } def get_known_periods(self) -> list: diff --git a/tests/core/test_disk_storage_delete_contained_periods.py b/tests/core/test_disk_storage_delete_contained_periods.py new file mode 100644 index 000000000..e7887b828 --- /dev/null +++ b/tests/core/test_disk_storage_delete_contained_periods.py @@ -0,0 +1,318 @@ +"""``OnDiskStorage.delete(period, branch_name)`` deletes the periods within +``period`` (policyengine-core#564). + +``Holder.delete_arrays`` documents that deleting a period removes "all values +for any period included in period", and ``InMemoryStorage.delete`` does that +(deleting ``2025`` removes ``2025-01`` ... ``2025-12``). ``OnDiskStorage.delete`` +used to remove only the file keyed exactly ``f"{branch_name}_{period}"``, so a +disk-backed simulation kept values the caller had deleted. + +Each example runs on both backends. ``test_disk_and_memory_delete_the_same_keys`` +checks every pair of a stored and a deleted period on a small grid, and +``test_disk_storage_delete_property.py`` checks random sequences of puts and +deletes. +""" + +import os +import warnings + +import numpy as np +import pytest + +from policyengine_core import periods +from policyengine_core.country_template import CountryTaxBenefitSystem +from policyengine_core.data_storage import InMemoryStorage, OnDiskStorage +from policyengine_core.experimental import MemoryConfig +from policyengine_core.simulations import SimulationBuilder + +BACKENDS = ("memory", "disk") + +# On Windows a ``:`` in a file name is invalid or names an alternate data +# stream, so disk storage cannot keep periods whose string form has one +# (``month:2025-01:3``, ``year:2025-03``). That predates this fix; tests of +# such periods run elsewhere. +skip_colon_periods_on_windows = pytest.mark.skipif( + os.name == "nt", + reason="On-disk keys for multi-unit periods contain ':', which Windows " + "file names cannot.", +) + + +@pytest.fixture +def make_storage(tmp_path): + def make(backend, is_eternal=False): + if backend == "memory": + return InMemoryStorage(is_eternal=is_eternal) + directory = tmp_path / f"storage_{len(os.listdir(tmp_path))}" + directory.mkdir() + return OnDiskStorage( + str(directory), is_eternal=is_eternal, preserve_storage_dir=True + ) + + return make + + +def put_all(storage, keys): + for index, (branch_name, period) in enumerate(keys): + storage.put(np.array([float(index)]), periods.period(period), branch_name) + + +def known(storage): + """The stored (branch name, period string) pairs, sorted.""" + return sorted( + (branch_name, str(period)) + for branch_name, period in storage.get_known_branch_periods() + ) + + +# ----- Examples on both backends ---------------------------------------------- + + +@pytest.mark.parametrize("backend", BACKENDS) +def test_delete_exact_period(make_storage, backend): + storage = make_storage(backend) + put_all(storage, [("default", "2025-01"), ("default", "2025-02")]) + + storage.delete(periods.period("2025-01"), "default") + + assert known(storage) == [("default", "2025-02")] + + +@pytest.mark.parametrize("backend", BACKENDS) +def test_delete_year_deletes_its_months_and_days(make_storage, backend): + storage = make_storage(backend) + put_all( + storage, + [ + ("default", "2024-12"), + ("default", "2025"), + ("default", "2025-01"), + ("default", "2025-12"), + ("default", "2025-03-15"), + ("default", "2026-01"), + ], + ) + + storage.delete(periods.period(2025), "default") + + assert known(storage) == [("default", "2024-12"), ("default", "2026-01")] + + +@pytest.mark.parametrize("backend", BACKENDS) +def test_delete_month_keeps_the_year_that_contains_it(make_storage, backend): + storage = make_storage(backend) + put_all( + storage, + [ + ("default", "2025"), + ("default", "2025-01"), + ("default", "2025-01-31"), + ("default", "2025-02"), + ("default", "2025-02-01"), + ], + ) + + storage.delete(periods.period("2025-01"), "default") + + assert known(storage) == [ + ("default", "2025"), + ("default", "2025-02"), + ("default", "2025-02-01"), + ] + + +@pytest.mark.parametrize("backend", BACKENDS) +def test_delete_keeps_other_branches(make_storage, backend): + storage = make_storage(backend) + put_all( + storage, + [ + ("default", "2025-01"), + ("reform", "2025-01"), + ("reform", "2025-02"), + ], + ) + + storage.delete(periods.period(2025), "reform") + + assert known(storage) == [("default", "2025-01")] + np.testing.assert_array_equal(storage.get("2025-01", "default"), [0.0]) + + +@pytest.mark.parametrize("backend", BACKENDS) +def test_delete_matches_whole_branch_names_containing_underscores( + make_storage, backend +): + # ``pre`` and ``pre_tcja`` are prefixes of ``pre_tcja_ctc``; the ``y`` + # branch's prefix ``y_`` starts the ``y_2019`` branch's keys. + branch_names = ("pre", "pre_tcja", "pre_tcja_ctc", "y", "y_2019", "trailing_") + storage = make_storage(backend) + put_all( + storage, + [ + (branch_name, month) + for branch_name in branch_names + for month in ("2025-01", "2025-02") + ], + ) + + storage.delete(periods.period(2025), "pre_tcja") + storage.delete(periods.period(2025), "y") + storage.delete(periods.period("2025-01"), "trailing_") + + assert known(storage) == [ + ("pre", "2025-01"), + ("pre", "2025-02"), + ("pre_tcja_ctc", "2025-01"), + ("pre_tcja_ctc", "2025-02"), + ("trailing_", "2025-02"), + ("y_2019", "2025-01"), + ("y_2019", "2025-02"), + ] + + +@pytest.mark.parametrize("backend", BACKENDS) +def test_eternal_storage_deletes_only_the_branch_eternity_key(make_storage, backend): + storage = make_storage(backend, is_eternal=True) + put_all(storage, [("default", "2025"), ("no_salt", "2025"), ("no", "2025")]) + + # Any period deletes the branch's single ETERNITY value. + storage.delete(periods.period("2025-01"), "no_salt") + + assert known(storage) == [("default", "ETERNITY"), ("no", "ETERNITY")] + + +@pytest.mark.parametrize("backend", BACKENDS) +def test_deleting_eternity_deletes_every_period_of_the_branch(make_storage, backend): + storage = make_storage(backend) + put_all( + storage, + [("default", "2024"), ("default", "2025-01"), ("reform", "2025-01")], + ) + + storage.delete(periods.period(periods.ETERNITY), "default") + + assert known(storage) == [("reform", "2025-01")] + + +@skip_colon_periods_on_windows +@pytest.mark.parametrize("backend", BACKENDS) +def test_delete_multi_unit_periods(make_storage, backend): + storage = make_storage(backend) + put_all( + storage, + [ + ("default", "month:2025-01:3"), + ("default", "year:2025-03"), + ("default", "2025-02"), + ("default", "2025-04"), + ], + ) + + # A quarter contains its months, but a month does not contain the + # quarter, and 2025 does not contain the year starting in March 2025. + storage.delete(periods.period("month:2025-01:3"), "default") + assert known(storage) == [("default", "2025-04"), ("default", "year:2025-03")] + + storage.delete(periods.period("2025-04"), "default") + assert known(storage) == [("default", "year:2025-03")] + + storage.delete(periods.period(2025), "default") + assert known(storage) == [("default", "year:2025-03")] + + storage.delete(periods.period("year:2025:2"), "default") + assert known(storage) == [] + + +# ----- Disk and memory agree on every pair of periods ------------------------- + +PERIODS = [ + "2024", + "2025", + "2025-01", + "2025-02", + "2025-12", + "2025-01-01", + "2025-01-31", + "2025-02-01", + "2026-01", + "ETERNITY", +] +if os.name != "nt": + PERIODS += ["month:2025-01:2", "year:2025-03", "year:2024:2"] +BRANCH_NAMES = ["default", "pre_tcja", "pre_tcja_ctc"] + + +@pytest.mark.parametrize("deleted_branch_name", BRANCH_NAMES) +@pytest.mark.parametrize("deleted_period", PERIODS) +def test_disk_and_memory_delete_the_same_keys( + make_storage, deleted_period, deleted_branch_name +): + keys = [(branch_name, period) for branch_name in BRANCH_NAMES for period in PERIODS] + memory, disk = make_storage("memory"), make_storage("disk") + put_all(memory, keys) + put_all(disk, keys) + deleted_period = periods.period(deleted_period) + + memory.delete(deleted_period, deleted_branch_name) + disk.delete(deleted_period, deleted_branch_name) + + assert known(disk) == known(memory) + # And both keep exactly the keys the period does not contain. + assert known(memory) == sorted( + (branch_name, str(periods.period(period))) + for branch_name, period in keys + if branch_name != deleted_branch_name + or not deleted_period.contains(periods.period(period)) + ) + for branch_name, period in known(memory): + np.testing.assert_array_equal( + disk.get(period, branch_name), memory.get(period, branch_name) + ) + + +# ----- Through the simulation (the #564 reproduction) ------------------------- + + +def simulation_with_rent(on_disk): + simulation = SimulationBuilder().build_from_entities( + CountryTaxBenefitSystem(), + {"persons": {"a": {}}, "households": {"h": {"parents": ["a"]}}}, + ) + if on_disk: + with warnings.catch_warnings(): + warnings.simplefilter("ignore") + simulation.memory_config = MemoryConfig(max_memory_occupation=0) + holder = simulation.get_holder("rent") + holder._disk_storage = holder.create_disk_storage() + holder._on_disk_storable = True + simulation.set_input("rent", "2025-01", [500.0]) + simulation.set_input("rent", "2025-02", [600.0]) + simulation.set_input("rent", "2026-01", [700.0]) + return simulation + + +@pytest.mark.parametrize("on_disk", [False, True], ids=["memory", "disk"]) +def test_simulation_delete_arrays_deletes_the_months_of_a_year(on_disk): + simulation = simulation_with_rent(on_disk) + holder = simulation.get_holder("rent") + if on_disk: + assert holder._memory_storage.get_known_periods() == [] + assert len(holder._disk_storage.get_known_periods()) == 3 + + simulation.delete_arrays("rent", "2025") + + assert holder.get_known_periods() == [periods.period("2026-01")] + assert holder.get_array("2025-01") is None + assert holder.get_array("2025-02") is None + np.testing.assert_array_equal(holder.get_array("2026-01"), [700.0]) + + +@pytest.mark.parametrize("on_disk", [False, True], ids=["memory", "disk"]) +def test_holder_delete_arrays_deletes_the_months_of_a_year(on_disk): + simulation = simulation_with_rent(on_disk) + holder = simulation.get_holder("rent") + + holder.delete_arrays("2025", "default") + + assert holder.get_known_periods() == [periods.period("2026-01")] diff --git a/tests/core/test_disk_storage_delete_property.py b/tests/core/test_disk_storage_delete_property.py new file mode 100644 index 000000000..571c4c3c8 --- /dev/null +++ b/tests/core/test_disk_storage_delete_property.py @@ -0,0 +1,121 @@ +"""Disk and memory storage keep the same values through any puts and deletes. + +Random sequences of ``put``, ``delete(period, branch_name)`` and +``delete(None, branch_name)`` run on an ``InMemoryStorage`` and an +``OnDiskStorage``, and on a reference model of the documented behaviour: +``delete(period, branch_name)`` removes the branch's values for every period +that ``period`` contains, ``delete(None, branch_name)`` removes all of the +branch's values, and eternal storage keeps one ``ETERNITY`` value per branch. +After every step both storages hold exactly the model's (branch, period) +keys, with the values last put there. + +Branch names include ones that contain ``_`` or start with another branch's +name followed by ``_``, since disk keys join branch and period with ``_``. +``test_disk_storage_delete_contained_periods.py`` pins the same behaviour +with examples. +""" + +import os +import tempfile + +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 import periods +from policyengine_core.data_storage import InMemoryStorage, OnDiskStorage + +BRANCH_NAMES = ( + "default", + "reform", + "pre", + "pre_tcja", + "pre_tcja_ctc", + "y", + "y_2019", + "trailing_", +) + + +@st.composite +def period_strategy(draw): + if draw(st.integers(0, 19)) == 0: + return periods.period(periods.ETERNITY) + unit = draw(st.sampled_from([periods.YEAR, periods.MONTH, periods.DAY])) + year = draw(st.integers(2023, 2026)) + if os.name == "nt": + # Windows file names cannot contain the ``:`` of a multi-unit or + # rolling-year period's string form. + month = 1 if unit == periods.YEAR else draw(st.integers(1, 12)) + size = 1 + else: + month = draw(st.integers(1, 12)) + size = draw(st.integers(1, 3)) + # In-memory storage rejects years and months anchored mid-month, whose + # string form drops the day. + day = draw(st.integers(1, 28)) if unit == periods.DAY else 1 + return periods.Period((unit, periods.Instant((year, month, day)), size)) + + +put_step = st.tuples( + st.just("put"), + period_strategy(), + st.sampled_from(BRANCH_NAMES), + st.integers(0, 1_000), +) +delete_step = st.tuples( + st.just("delete"), + st.one_of(st.none(), period_strategy()), + st.sampled_from(BRANCH_NAMES), + st.none(), +) + + +def stored(storage): + """Map each stored (branch name, period string) to its value.""" + return { + (branch_name, str(period)): float(storage.get(period, branch_name)[0]) + for branch_name, period in storage.get_known_branch_periods() + } + + +@hypothesis.settings( + max_examples=200, + deadline=None, + suppress_health_check=[hypothesis.HealthCheck.too_slow], +) +@hypothesis.given( + is_eternal=st.booleans(), + steps=st.lists(st.one_of(put_step, put_step, delete_step), max_size=25), +) +def test_disk_and_memory_storage_delete_the_same_keys(is_eternal, steps): + model = {} + memory = InMemoryStorage(is_eternal=is_eternal) + with tempfile.TemporaryDirectory() as directory: + disk = OnDiskStorage( + directory, is_eternal=is_eternal, preserve_storage_dir=True + ) + for action, period, branch_name, value in steps: + if period is not None and is_eternal: + period = periods.period(periods.ETERNITY) + if action == "put": + memory.put(np.array([float(value)]), period, branch_name) + disk.put(np.array([float(value)]), period, branch_name) + model[(branch_name, str(period))] = float(value) + else: + memory.delete(period, branch_name) + disk.delete(period, branch_name) + model = { + (key_branch_name, key_period): key_value + for (key_branch_name, key_period), key_value in model.items() + if key_branch_name != branch_name + or ( + period is not None + and not period.contains(periods.period(key_period)) + ) + } + + assert stored(disk) == stored(memory) == model From a96bc647487e21f2f1151aca4083647077d4c462 Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Tue, 6 Oct 2026 10:06:19 -0400 Subject: [PATCH 04/15] Use master's readable-branch lookup for branch-scoped known periods Reconcile #552 with master after the merge: - Holder.get_known_periods(branch_name) uses master's _readable_branches (added by #562) instead of a second copy of the lookup, and lists each period once. get_array is master's again. - _dump_holder reads the dumped simulation's branch from the holder, as #560 does, and keeps master's derived_periods.txt with each value's mark read on that branch. - The _calculate comment says what scoping still changes now that uprating and carry-over read only inputs (#562/#563): a later period stored under an unreadable branch no longer leaves a default uncached. Tests: the isolation grid now also compares what the branch then reads (periods and derived marks); a period stored under several readable branches is listed once; a restored branch dump calculates what the branch calculates; and a Hypothesis differential checks get_known_periods against get_array and get_input_periods against is_derived for random values on a five-branch lineage, in memory and on disk. Co-Authored-By: Claude Opus 5.5 --- policyengine_core/holders/holder.py | 90 +++++++------- policyengine_core/simulations/simulation.py | 7 +- policyengine_core/tools/simulation_dumper.py | 18 +-- .../test_known_periods_branch_visibility.py | 91 ++++++++++++-- ...nown_periods_branch_visibility_property.py | 111 ++++++++++++++++++ 5 files changed, 244 insertions(+), 73 deletions(-) create mode 100644 tests/core/test_known_periods_branch_visibility_property.py diff --git a/policyengine_core/holders/holder.py b/policyengine_core/holders/holder.py index 10332e325..54b7ae129 100644 --- a/policyengine_core/holders/holder.py +++ b/policyengine_core/holders/holder.py @@ -118,53 +118,45 @@ def _get_array_from_storage( value = self._disk_storage.get(period, branch_name) return value - def _readable_branch_names(self, branch_name: str = "default") -> List[str]: - """ - Branches whose stored values ``get_array(period, branch_name)`` can - return, in lookup order. - - That is ``branch_name`` itself, then (unless it is ``default``) each - ``simulation.parent_branch`` ancestor, then ``default``. Nested - branches inherit values from their parent (e.g. a ``no_salt`` branch - cloned from an ``itemizing`` branch still sees ``tax_unit_itemizes`` - set on the ``itemizing`` branch). Previously the fallback returned - the first branch in dict-insertion order (bug C1) — silently swapping - values between unrelated sibling branches (reform vs baseline) and - producing wrong reform deltas. The post-C1 behavior only fell back to - ``default``, which broke country-package nested-branch patterns that - relied on the ancestor's input being visible. - """ - if branch_name == "default": - return ["default"] - branch_names = [branch_name] - simulation = getattr(self, "simulation", None) - parent = getattr(simulation, "parent_branch", None) if simulation else None - while parent is not None: - branch_names.append(parent.branch_name) - parent = getattr(parent, "parent_branch", None) - branch_names.append("default") - return list(dict.fromkeys(branch_names)) - def get_array(self, period: Period, branch_name: str = "default") -> ArrayLike: """ Get the value of the variable for the given period. - Values stored under ``branch_name``, its ``parent_branch`` ancestors - and ``default`` are visible, in that order. If the value is not - known on any of them, return ``None``. + If the value is not known, return ``None``. """ if self.variable.is_neutralized: return self.default_array() - # The branch's own value is the common case: look it up before - # walking the ancestors. value = self._get_array_from_storage(period, branch_name) if value is not None: return value - for readable_branch_name in self._readable_branch_names(branch_name)[1:]: - value = self._get_array_from_storage(period, readable_branch_name) - if value is not None: - return value - return None + if value is None and branch_name != "default": + # Walk up ``simulation.parent_branch`` so nested branches inherit + # values from their parent (e.g. a ``no_salt`` branch cloned + # from an ``itemizing`` branch still sees ``tax_unit_itemizes`` + # set on the ``itemizing`` branch). Fall back to ``default`` + # only if no ancestor branch has a value. Previously the + # fallback returned the first branch in dict-insertion order + # (bug C1) — silently swapping values between unrelated + # sibling branches (reform vs baseline) and producing wrong + # reform deltas. The post-C1 behavior only fell back to + # ``default``, which broke country-package nested-branch + # patterns that relied on the ancestor's input being visible. + parent = ( + getattr(self.simulation, "parent_branch", None) + if self.simulation + else None + ) + while parent is not None: + ancestor_value = self._get_array_from_storage( + period, + parent.branch_name, + ) + if ancestor_value is not None: + return ancestor_value + parent = getattr(parent, "parent_branch", None) + default_value = self._get_array_from_storage(period, "default") + if default_value is not None: + return default_value def get_memory_usage(self) -> dict: """ @@ -217,19 +209,21 @@ def get_known_periods(self, branch_name: str = None) -> List[Period]: """ Get the list of periods the variable value is known for. - With ``branch_name``, list only the periods ``get_array(period, - branch_name)`` can read: those stored under that branch, its - ``parent_branch`` ancestors or ``default``. Without it, list the - periods stored under every branch in this holder, some of which that - branch may not be able to read. + With ``branch_name``, list each period ``get_array(period, + branch_name)`` reads a stored value for, once: those stored under + that branch, its ``parent_branch`` ancestors or ``default`` (see + ``_readable_branches``). Without it, list the period of every value + stored, under any branch; that branch may not read some of them. """ if branch_name is not None: - readable_branch_names = set(self._readable_branch_names(branch_name)) - return [ - period - for stored_branch_name, period in self.get_known_branch_periods() - if stored_branch_name in readable_branch_names - ] + readable = set(self._readable_branches(branch_name)) + return list( + dict.fromkeys( + period + for stored_branch, period in self.get_known_branch_periods() + if stored_branch in readable + ) + ) return list(self._memory_storage.get_known_periods()) + list( (self._disk_storage.get_known_periods() if self._disk_storage else []) ) diff --git a/policyengine_core/simulations/simulation.py b/policyengine_core/simulations/simulation.py index 2b8624cb8..0c708511d 100644 --- a/policyengine_core/simulations/simulation.py +++ b/policyengine_core/simulations/simulation.py @@ -1083,9 +1083,10 @@ def _calculate(self, variable_name: str, period: Period = None) -> ArrayLike: # If no result, use the default value and cache it if array is None: # Check if the variable has a previously defined value. - # Only periods this branch can read count: a period stored - # only under an unrelated branch would read back as ``None`` - # and reach the arithmetic below. + # Only periods this branch can read count: a value stored + # only under a branch this one cannot read must not change + # what it calculates or caches (below, a later stored period + # leaves the default uncached). known_periods = holder.get_known_periods(self.branch_name) earlier_known_periods = [ known_period diff --git a/policyengine_core/tools/simulation_dumper.py b/policyengine_core/tools/simulation_dumper.py index a3b477fc1..da665e863 100644 --- a/policyengine_core/tools/simulation_dumper.py +++ b/policyengine_core/tools/simulation_dumper.py @@ -36,7 +36,7 @@ def dump_simulation(simulation, directory): # Dump variable values for holder in entity._holders.values(): - _dump_holder(holder, directory, simulation.branch_name) + _dump_holder(holder, directory) def restore_simulation(directory, tax_benefit_system, **kwargs): @@ -68,16 +68,16 @@ def restore_simulation(directory, tax_benefit_system, **kwargs): return simulation -def _dump_holder(holder, directory, branch_name="default"): - # Dump the values the simulation's branch reads. A holder can also store - # periods under branches this one cannot see; reading those back under - # this branch gives ``None``, which would be saved as an object array - # that ``restore_simulation`` cannot load. +def _dump_holder(holder, directory): disk_storage = holder.create_disk_storage(directory, preserve=True) + # Dump what the simulation reads: on a branch, its own value, else its + # nearest ancestor's or the default one. A value stored only under a + # branch the simulation cannot read is left out: read here, it would be + # ``None``, saved as an array ``restore_simulation`` cannot load. + branch_name = holder.simulation.branch_name derived_periods = set() - for period in dict.fromkeys(holder.get_known_periods(branch_name)): - value = holder.get_array(period, branch_name) - disk_storage.put(value, period) + for period in holder.get_known_periods(branch_name): + disk_storage.put(holder.get_array(period, branch_name), period) # Read the mark of exactly the value dumped: the same period on the # same branch as ``get_array``. if holder.is_derived(period, branch_name): diff --git a/tests/core/test_known_periods_branch_visibility.py b/tests/core/test_known_periods_branch_visibility.py index d1806af9b..00835c3e2 100644 --- a/tests/core/test_known_periods_branch_visibility.py +++ b/tests/core/test_known_periods_branch_visibility.py @@ -4,13 +4,18 @@ ``"no_salt_2018"`` on disk), and a holder can hold keys for branches other than the simulation's own. ``Holder.get_known_periods()`` lists the periods of every key, while ``Holder.get_array(period, branch_name)`` reads only that -branch, its ``parent_branch`` ancestors and ``default``. Three defects came -from that gap: - -* ``Simulation._calculate`` took the latest known earlier period from the - unscoped list. If that period was stored only under a branch the simulation - cannot read, ``get_array`` returned ``None``: uprating raised ``TypeError`` - (``None * factor``) and auto-carry-over cached ``NaN``. +branch, its ``parent_branch`` ancestors and ``default``. +``Holder.get_known_periods(branch_name)`` lists exactly the periods +``get_array`` reads for that branch, each once. Three defects came from the +gap between the two: + +* ``Simulation._calculate`` checked the unscoped list for stored periods. When + it took the latest earlier period from it, a period stored only under a + branch the simulation cannot read gave ``None``: uprating raised + ``TypeError`` (``None * factor``) and auto-carry-over cached ``NaN``. Since + both read only inputs the branch reads (``Holder.get_input_periods``), such + a period still made the simulation leave a default uncached that it would + otherwise cache. * ``OnDiskStorage`` parsed ``f"{branch}_{period}"`` keys with ``split("_")[1]``. A branch name containing ``_`` raised ``ValueError`` (``no_salt`` -> period ``"salt"``) or listed the wrong period @@ -18,7 +23,11 @@ ``pre_tcja_ctc``, whose key shares the prefix. * ``dump_simulation`` read every listed period under ``default``, saving ``None`` for a period stored only on the dumped branch, which - ``restore_simulation`` could not load. + ``restore_simulation`` could not load, and the derived marks of the + ``default`` branch's values rather than those of the values dumped. + +``test_known_periods_branch_visibility_property.py`` checks the listing +against ``get_array`` and ``get_input_periods`` for random stored values. """ import itertools @@ -181,9 +190,9 @@ def test_carry_over_ignores_later_period_on_unrelated_branch(carry_over_system): # Invariant: values stored only under branches a simulation cannot read do not -# change what it calculates. Checked exhaustively against a simulation that -# never had those values, for every combination of readable years, unreadable -# years and requested year below. +# change what it calculates, nor what it stores where it reads. Checked +# exhaustively against a simulation that never had those values, for every +# combination of readable years, unreadable years and requested year below. READABLE_YEAR_SETS = [(), (2016,), (2018,), (2016, 2018)] UNREADABLE_YEAR_SETS = [ years @@ -194,13 +203,20 @@ def test_carry_over_ignores_later_period_on_unrelated_branch(carry_over_system): def _calculate_with(system, variable, readable, unreadable, requested): + """The branch's result, and the (period, derived) pairs it then reads.""" simulation = new_simulation(system) branch = simulation.get_branch("reform") for year in readable: store(branch, variable, year, 1_000.0 * (year - 2000), "reform") for year in unreadable: store(branch, variable, year, 7_777.0, "baseline") - return only(branch.calculate(variable, requested)) + result = only(branch.calculate(variable, requested)) + holder = branch.get_holder(variable) + read = sorted( + (str(period), holder.is_derived(period, "reform")) + for period in holder.get_known_periods("reform") + ) + return result, read @pytest.mark.parametrize("requested", REQUESTED_YEARS) @@ -228,7 +244,7 @@ def test_unreadable_branch_periods_do_not_change_results( tax_benefit_system, variable, readable, (), requested ) - assert math.isfinite(with_unreadable) + assert math.isfinite(with_unreadable[0]) assert with_unreadable == without_unreadable @@ -292,6 +308,24 @@ def test_known_periods_for_branch_are_exactly_the_readable_ones( assert len(listed) == len(set(listed)), reader_name +@pytest.mark.parametrize("on_disk", [False, True]) +def test_known_periods_for_branch_list_each_period_once(system, tmp_path, on_disk): + """A period stored under several branches a reader reads, or both in + memory and on disk, is listed once.""" + period = periods.period(2016) + for reader_name, reader in _lineage(system).items(): + holder = reader.get_holder("uprated_income") + directory = tmp_path / reader_name.replace(" ", "-") + directory.mkdir() + holder._disk_storage = holder.create_disk_storage(str(directory), preserve=True) + storages = [holder._memory_storage, holder._disk_storage] + for branch_name in STORED_BRANCHES: + for storage in storages if on_disk else storages[:1]: + storage.put(np.array([1.0]), period, branch_name) + + assert holder.get_known_periods(reader.branch_name) == [period], reader_name + + def test_get_known_periods_without_branch_lists_every_branch(system): simulation = new_simulation(system) store(simulation, "uprated_income", 2016, 1.0, "default") @@ -426,3 +460,34 @@ def test_dump_branch_saves_the_values_the_branch_reads(system): assert only(holder.get_array(2016)) == 1_000 assert only(holder.get_array(2017)) == 3_000 assert only(holder.get_array(2018)) == 5_000 + + +def test_restored_branch_dump_calculates_what_the_branch_calculates(system): + """The dump keeps the derived mark of each value the branch reads, so the + restored simulation uprates from the same inputs as the branch.""" + simulation = new_simulation(system) + simulation.set_input("uprated_income", 2016, [1_000.0]) + # Calculated on ``default``, then set as an input on the branch. + simulation.calculate("uprated_income", 2019) + branch = simulation.get_branch("reform") + branch.set_input("uprated_income", 2019, [3_000.0]) + # Calculated on the branch only. + branch.calculate("uprated_income", 2017) + + directory = os.path.join(tempfile.mkdtemp(), "dump") + simulation_dumper.dump_simulation(branch, directory) + restored = simulation_dumper.restore_simulation(directory, system) + + holder = restored.get_holder("uprated_income") + assert sorted( + (str(period), holder.is_derived(period)) + for period in holder.get_known_periods() + ) == [("2016", False), ("2017", True), ("2019", False)] + for year in range(2016, 2023): + assert only(restored.calculate("uprated_income", year)) == pytest.approx( + only(branch.calculate("uprated_income", year)) + ), year + assert only(restored.calculate("uprated_income", 2021)) == pytest.approx( + 3_000 * growth(2019, 2021) + ) + diff --git a/tests/core/test_known_periods_branch_visibility_property.py b/tests/core/test_known_periods_branch_visibility_property.py new file mode 100644 index 000000000..e2f828ee3 --- /dev/null +++ b/tests/core/test_known_periods_branch_visibility_property.py @@ -0,0 +1,111 @@ +"""``Holder.get_known_periods(branch_name)`` agrees with ``get_array`` for any +stored values. + +Random values are stored, in memory or on disk, as inputs or derived, under +the branches of one lineage: ``default``, its branch ``a``, ``a``'s branch +``a_b``, ``default``'s other branch ``c``, and a branch named ``default`` +nested under ``a``. Each branch of the lineage, reading its own holder, sees: + +* ``get_known_periods(branch_name)`` lists each period ``get_array(period, + branch_name)`` reads a stored value for, once, and no other period: the + listing against the per-period reference, ``get_array``; +* ``get_input_periods(branch_name)`` lists the periods it lists whose value + is an input (``not is_derived(period, branch_name)``): the one pass over + the stored keys against the per-period reference. + +``test_known_periods_branch_visibility.py`` pins the same with examples. +""" + +import os +import tempfile + +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 import periods +from policyengine_core.country_template import CountryTaxBenefitSystem +from policyengine_core.country_template.entities import Person +from policyengine_core.model_api import YEAR, Variable +from policyengine_core.simulations import SimulationBuilder + +STORED_BRANCHES = ("default", "a", "a_b", "c", "unrelated") +YEARS = (2015, 2016, 2017) + + +class yearly_income(Variable): + value_type = float + entity = Person + definition_period = YEAR + label = "Yearly income" + + +@pytest.fixture(scope="module") +def system(): + system = CountryTaxBenefitSystem() + system.add_variable(yearly_income) + return system + + +def lineage(system): + simulation = SimulationBuilder().build_default_simulation(system, count=1) + a = simulation.get_branch("a") + return [ + simulation, + a, + a.get_branch("a_b"), + simulation.get_branch("c"), + a.get_branch("default"), + ] + + +stored_value = st.tuples( + st.sampled_from(STORED_BRANCHES), + st.sampled_from(YEARS), + st.booleans(), # derived + st.booleans(), # on disk +) + + +@hypothesis.settings( + max_examples=100, + deadline=None, + suppress_health_check=[hypothesis.HealthCheck.too_slow], +) +@hypothesis.given(stored=st.lists(stored_value, max_size=12)) +def test_known_periods_for_a_branch_are_the_periods_it_reads(system, stored): + with tempfile.TemporaryDirectory() as directory: + for index, reader in enumerate(lineage(system)): + holder = reader.get_holder("yearly_income") + reader_directory = os.path.join(directory, str(index)) + os.mkdir(reader_directory) + holder._disk_storage = holder.create_disk_storage( + reader_directory, preserve=True + ) + for value, (branch_name, year, derived, on_disk) in enumerate(stored): + storage = holder._disk_storage if on_disk else holder._memory_storage + storage.put( + np.array([float(value)]), + periods.period(year), + branch_name, + derived=derived, + ) + + branch_name = reader.branch_name + listed = holder.get_known_periods(branch_name) + read = [ + period + for period in map(periods.period, YEARS) + if holder.get_array(period, branch_name) is not None + ] + + assert len(listed) == len(set(listed)) + assert set(listed) == set(read) + assert set(holder.get_input_periods(branch_name)) == { + period + for period in listed + if not holder.is_derived(period, branch_name) + } From db6c3f43a3851d0158e75bbb772a51818dfcb92a Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Wed, 7 Oct 2026 17:00:48 -0400 Subject: [PATCH 05/15] Format the branch visibility regression grid Co-Authored-By: Claude Opus 5.5 --- tests/core/test_known_periods_branch_visibility.py | 1 - 1 file changed, 1 deletion(-) diff --git a/tests/core/test_known_periods_branch_visibility.py b/tests/core/test_known_periods_branch_visibility.py index 00835c3e2..df7ea35c1 100644 --- a/tests/core/test_known_periods_branch_visibility.py +++ b/tests/core/test_known_periods_branch_visibility.py @@ -490,4 +490,3 @@ def test_restored_branch_dump_calculates_what_the_branch_calculates(system): assert only(restored.calculate("uprated_income", 2021)) == pytest.approx( 3_000 * growth(2019, 2021) ) - From 446a1ba936fef796f912c7bf192c980536366346 Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Tue, 6 Oct 2026 10:08:15 -0400 Subject: [PATCH 06/15] Reject on disk the keys memory rejects, and restore only keys Fixes the two findings of #565's review (REQUEST_CHANGES at bfa78cba): 1. An input for a month or year anchored mid-month (month:2025-03-15) was accepted on disk under the calendar month's key, since its string form drops the day; deleting its own period, which does not contain that calendar month, then left it. In-memory storage already rejected it. Both storages now run one check, storage_keys.check_storable, so whether an input is accepted no longer depends on which storage Holder._set picks by memory use; disk storage rejects branch names containing ':' too. 2. restore() read back every .npy file. A name that is not a key put writes (stray.npy, default_garbage.npy, default_2025-3.npy) made delete and get_known_periods raise, so a valid deletion left its value. restore now leaves such files out, with a warning naming them. Tests: examples on both backends and through Simulation for each finding; the memory/disk Hypothesis differential now draws mid-month periods and a ':' branch name and checks both storages accept or reject each put alike; a new Hypothesis property restores a directory of random keys and non-key names and checks restore reads back exactly the keys, warns about the rest, and then deletes like memory. Folds #565's changelog fragment into this PR's. Co-Authored-By: Claude Opus 5.5 --- ...fix-disk-delete-contained-periods.fixed.md | 1 - .../fix-uprating-branch-visibility.fixed.md | 2 +- docs/python_api/data_storage.md | 7 + .../data_storage/in_memory_storage.py | 17 +-- .../data_storage/on_disk_storage.py | 38 ++++++ .../data_storage/storage_keys.py | 35 +++++ ...t_disk_storage_delete_contained_periods.py | 124 +++++++++++++++++ .../core/test_disk_storage_delete_property.py | 127 ++++++++++++++++-- .../test_known_periods_branch_visibility.py | 15 +++ 9 files changed, 339 insertions(+), 27 deletions(-) delete mode 100644 changelog.d/fix-disk-delete-contained-periods.fixed.md create mode 100644 policyengine_core/data_storage/storage_keys.py diff --git a/changelog.d/fix-disk-delete-contained-periods.fixed.md b/changelog.d/fix-disk-delete-contained-periods.fixed.md deleted file mode 100644 index 9748a6f83..000000000 --- a/changelog.d/fix-disk-delete-contained-periods.fixed.md +++ /dev/null @@ -1 +0,0 @@ -On-disk storage now deletes every period the deleted period contains, as in-memory storage does, so deleting `2025` from a disk-backed holder also deletes `2025-01` through `2025-12`. diff --git a/changelog.d/fix-uprating-branch-visibility.fixed.md b/changelog.d/fix-uprating-branch-visibility.fixed.md index 36eb6f762..708932a71 100644 --- a/changelog.d/fix-uprating-branch-visibility.fixed.md +++ b/changelog.d/fix-uprating-branch-visibility.fixed.md @@ -1 +1 @@ -Uprating and auto-carry-over now use only periods the current branch can read (its own, its parent branches' and `default`), so a period stored only under an unrelated branch no longer raises `TypeError` or caches `NaN`; on-disk storage parses branch names that contain `_`, and `dump_simulation` saves the values the dumped branch reads. +`Holder.get_known_periods(branch_name)` lists only the periods that branch reads, and `Simulation` uses it, so values stored only under other branches no longer change what a branch calculates or caches; on-disk storage now parses branch names containing `_`, deletes the periods within a deleted period as in-memory storage does, rejects the keys in-memory storage rejects (years and months anchored mid-month, branch names containing `:`) and leaves files that are not its keys out of `restore`; and `dump_simulation` saves the values the dumped branch reads, with their derived marks. diff --git a/docs/python_api/data_storage.md b/docs/python_api/data_storage.md index 4307934f6..2227f4999 100644 --- a/docs/python_api/data_storage.md +++ b/docs/python_api/data_storage.md @@ -20,6 +20,13 @@ The `policyengine_core.data_storage` module contains the classes that are used t :show-inheritance: ``` +## Storage keys + +```{eval-rst} +.. automodule:: policyengine_core.data_storage.storage_keys + :members: +``` + ## TemporaryStorageDirectory ```{eval-rst} diff --git a/policyengine_core/data_storage/in_memory_storage.py b/policyengine_core/data_storage/in_memory_storage.py index f1c3304f4..231440eb3 100644 --- a/policyengine_core/data_storage/in_memory_storage.py +++ b/policyengine_core/data_storage/in_memory_storage.py @@ -4,6 +4,7 @@ from numpy.typing import ArrayLike from policyengine_core import periods +from policyengine_core.data_storage.storage_keys import check_storable from policyengine_core.periods import Period @@ -266,20 +267,8 @@ def put( period = periods.period(period) # Keys embed f"{branch_name}:{period}", so reject inputs whose key - # would be ambiguous to parse back: branch names containing the - # separator, and month/year periods anchored mid-month (their string - # form drops the day). See policyengine-core#526 for the structured- - # key refactor that will lift these restrictions. - if ":" in branch_name: - raise ValueError( - f"Branch name {branch_name!r} may not contain ':' " - "(see policyengine-core#526)." - ) - if period.unit in (periods.YEAR, periods.MONTH) and period.start.day != 1: - raise ValueError( - f"Cannot cache period {period} anchored mid-month: its " - "string form is lossy (see policyengine-core#526)." - ) + # would be ambiguous to parse back, as ``OnDiskStorage.put`` does. + check_storable(branch_name, period) key = f"{branch_name}:{period}" self._arrays[key] = value self._stop_sharing(key) diff --git a/policyengine_core/data_storage/on_disk_storage.py b/policyengine_core/data_storage/on_disk_storage.py index b9a1cc5f7..1dcdfc328 100644 --- a/policyengine_core/data_storage/on_disk_storage.py +++ b/policyengine_core/data_storage/on_disk_storage.py @@ -1,6 +1,7 @@ import os import shutil import tempfile +import warnings import weakref import numpy @@ -11,6 +12,7 @@ directory_containing, path_key, ) +from policyengine_core.data_storage.storage_keys import check_storable from policyengine_core.enums import EnumArray from policyengine_core.periods import Period @@ -59,6 +61,19 @@ def _split_key(key: str) -> tuple: return branch_name, period +def _is_key(key: str) -> bool: + """Whether ``key`` is one ``put`` writes: ``f"{branch_name}_{period}"``, + with the period in its string form, so that it splits back into the + branch name and the period it was written for.""" + _, separator, period = key.rpartition("_") + if not separator: + return False + try: + return str(periods.period(period)) == period + except ValueError: + return False + + def _is_within(key: str, branch_name: str, period: Period) -> bool: """Whether file key ``key`` stores a value of ``branch_name`` for a period that ``period`` contains.""" @@ -295,6 +310,11 @@ def put( if self.is_eternal: period = periods.period(periods.ETERNITY) period = periods.period(period) + # The key must split back into this branch name and period (see + # ``get_known_branch_periods`` and ``delete``): reject what + # ``InMemoryStorage.put`` rejects. A month anchored mid-month, say, + # would be stored under the key of the calendar month. + check_storable(branch_name, period) filename = f"{branch_name}_{period}" path = self._path_to_write(filename) @@ -388,6 +408,10 @@ def restore(self) -> None: """Read back the values stored in this storage's directory: for each key, the file named for it (not those in ``REPLACEMENTS_DIR``). + A ``.npy`` file whose name is not a key ``put`` writes (a branch + name, ``_`` and a period in its string form) holds no value of this + storage: it is left out, with a warning. + This storage writes over those files, as over those it wrote, until it is cloned or copied, or this process forks: from then on it writes a new file instead, since the clone, copy or forked process reads @@ -399,14 +423,28 @@ def restore(self) -> None: self._files = files = {} # Files read back from a directory carry no derived marks. self._derived = set() + not_keys = [] # Restore self._files from content of storage_dir. for filename in os.listdir(self.storage_dir): if not filename.endswith(".npy"): continue path = os.path.join(self.storage_dir, filename) filename_core = filename.rsplit(".", 1)[0] + if not _is_key(filename_core): + # Reading it back would put a key in ``_files`` that names no + # branch and period, which ``delete`` and + # ``get_known_periods`` could not parse. + not_keys.append(filename) + continue files[filename_core] = path self._restored_paths.update(files.values()) + if not_keys: + warnings.warn( + f"Not restoring {', '.join(sorted(not_keys))} in " + f"{self.storage_dir}: a stored value's file is named for a " + "branch name, '_' and a period.", + stacklevel=2, + ) def __del__(self, _rmtree=shutil.rmtree, _getpid=os.getpid) -> None: # (The defaults keep the two functions reachable while the diff --git a/policyengine_core/data_storage/storage_keys.py b/policyengine_core/data_storage/storage_keys.py new file mode 100644 index 000000000..bb5575026 --- /dev/null +++ b/policyengine_core/data_storage/storage_keys.py @@ -0,0 +1,35 @@ +"""What a storage key can hold. + +Both storages key a value by its branch name and its period's string form: +``f"{branch_name}:{period}"`` in memory, ``f"{branch_name}_{period}"`` on disk +(the name of the value's file). ``check_storable`` rejects what such a key +cannot hold, the same way for both storages, so whether a holder accepts a +value never depends on which of the two it goes to (``Holder._set`` chooses by +how much memory is in use). +""" + +from policyengine_core import periods +from policyengine_core.periods import Period + + +def check_storable(branch_name: str, period: Period) -> None: + """Raise ``ValueError`` if no storage key can hold ``branch_name`` and + ``period``. + + Two cases: a branch name containing ``:``, which separates the parts of an + in-memory key, and a year or month anchored mid-month, whose string form + drops the day (``month:2025-03-15`` is ``2025-03``), so that its key would + be the key of another period. See policyengine-core#526 for the + structured-key refactor that will lift these restrictions. + """ + if ":" in branch_name: + raise ValueError( + f"Branch name {branch_name!r} may not contain ':' " + "(see policyengine-core#526)." + ) + if period.unit in (periods.YEAR, periods.MONTH) and period.start.day != 1: + raise ValueError( + f"Cannot cache period {period} anchored mid-month (it starts on " + f"{period.start}): its string form is lossy (see " + "policyengine-core#526)." + ) diff --git a/tests/core/test_disk_storage_delete_contained_periods.py b/tests/core/test_disk_storage_delete_contained_periods.py index e7887b828..cad439f5f 100644 --- a/tests/core/test_disk_storage_delete_contained_periods.py +++ b/tests/core/test_disk_storage_delete_contained_periods.py @@ -316,3 +316,127 @@ def test_holder_delete_arrays_deletes_the_months_of_a_year(on_disk): holder.delete_arrays("2025", "default") assert holder.get_known_periods() == [periods.period("2026-01")] + + +# ----- Keys no storage can hold (#565 review, finding 1) ---------------------- + +# A year or month anchored mid-month has a string form that drops the day +# (``month:2025-03-15`` is ``2025-03``), and ``:`` separates the parts of an +# in-memory key. In-memory storage rejected both (policyengine-core#526); disk +# storage stored a mid-month month under the calendar month's key, so reading +# the calendar month returned it, and deleting the period it was put for, +# which does not contain that calendar month, left it. +UNSTORABLE_KEYS = [ + ("default", periods.Period((periods.MONTH, periods.Instant((2025, 3, 15)), 1))), + ("default", periods.Period((periods.MONTH, periods.Instant((2025, 3, 2)), 2))), + ("default", periods.Period((periods.YEAR, periods.Instant((2025, 3, 15)), 1))), + ("no_salt", periods.Period((periods.YEAR, periods.Instant((2025, 1, 2)), 2))), + ("my:reform", periods.period(2025)), +] + + +@pytest.mark.parametrize("backend", BACKENDS) +@pytest.mark.parametrize( + "branch_name, period", + UNSTORABLE_KEYS, + ids=[ + f"{branch}-{period.unit}-{period.start}" for branch, period in UNSTORABLE_KEYS + ], +) +def test_put_rejects_keys_no_storage_can_hold( + make_storage, backend, branch_name, period +): + storage = make_storage(backend) + put_all(storage, [("default", "2025-03")]) + + with pytest.raises(ValueError, match="policyengine-core#526"): + storage.put(np.array([9.0]), period, branch_name) + + assert known(storage) == [("default", "2025-03")] + np.testing.assert_array_equal(storage.get("2025-03", "default"), [0.0]) + + +@pytest.mark.parametrize("on_disk", [False, True], ids=["memory", "disk"]) +def test_mid_month_input_is_rejected_whichever_storage_it_goes_to(on_disk): + # The review's reproduction: on disk, this input was accepted, then + # ``delete_arrays`` of its own period left it ([500.] instead of None). + simulation = simulation_with_rent(on_disk) + holder = simulation.get_holder("rent") + + with pytest.raises(ValueError, match="anchored mid-month"): + simulation.set_input("rent", "month:2025-03-15", [500.0]) + simulation.delete_arrays("rent", "month:2025-03-15") + + assert holder.get_array("month:2025-03-15") is None + assert holder.get_array("2025-03") is None + assert sorted(map(str, holder.get_known_periods())) == [ + "2025-01", + "2025-02", + "2026-01", + ] + + +# ----- Files restored from a directory (#565 review, finding 2) --------------- + +# ``restore`` read back every ``.npy`` file in the directory. One whose name +# is not a key ``put`` writes (no ``_``, no period after the last ``_``, or a +# period not in its string form) names no value, and deleting any period of +# any branch then raised parsing it, leaving the value it was to delete. +# Names without ``:``, which Windows file names cannot hold. +NOT_KEYS = [ + "stray.npy", + "default_garbage.npy", + "default_.npy", + "_.npy", + "default_2025-3.npy", + "default_2025-03-1.npy", +] +if os.name != "nt": + NOT_KEYS += ["default_month:2025-01:12.npy", "no_salt_year:2025:1.npy"] + + +def test_restore_leaves_out_files_that_are_not_keys(tmp_path): + storage = OnDiskStorage(str(tmp_path), preserve_storage_dir=True) + put_all(storage, [("default", "2025-01"), ("no_salt", "2025-02")]) + for name in NOT_KEYS: + np.save(str(tmp_path / name), np.array([9.0])) + + restored = OnDiskStorage(str(tmp_path), preserve_storage_dir=True) + with pytest.warns(UserWarning, match="Not restoring") as caught: + restored.restore() + + assert known(restored) == [("default", "2025-01"), ("no_salt", "2025-02")] + (warning,) = caught + for name in NOT_KEYS: + assert name in str(warning.message) + + # Deleting works as before the files were added. + restored.delete(periods.period(2025), "default") + assert restored.get("2025-01", "default") is None + assert known(restored) == [("no_salt", "2025-02")] + restored.delete(None, "no_salt") + assert known(restored) == [] + + +def test_restore_reads_back_every_key_without_warning(tmp_path): + keys = [ + ("default", "2025"), + ("no_salt", "2025-03"), + ("y_2019", "2025-03-05"), + ("trailing_", "ETERNITY"), + ] + if os.name != "nt": + keys += [("pre_tcja", "year:2025-03"), ("default", "month:2025-01:3")] + storage = OnDiskStorage(str(tmp_path), preserve_storage_dir=True) + put_all(storage, keys) + + restored = OnDiskStorage(str(tmp_path), preserve_storage_dir=True) + with warnings.catch_warnings(): + warnings.simplefilter("error") + restored.restore() + + assert known(restored) == known(storage) + for branch_name, period in keys: + np.testing.assert_array_equal( + restored.get(period, branch_name), storage.get(period, branch_name) + ) diff --git a/tests/core/test_disk_storage_delete_property.py b/tests/core/test_disk_storage_delete_property.py index 571c4c3c8..f384ec948 100644 --- a/tests/core/test_disk_storage_delete_property.py +++ b/tests/core/test_disk_storage_delete_property.py @@ -3,11 +3,20 @@ Random sequences of ``put``, ``delete(period, branch_name)`` and ``delete(None, branch_name)`` run on an ``InMemoryStorage`` and an ``OnDiskStorage``, and on a reference model of the documented behaviour: -``delete(period, branch_name)`` removes the branch's values for every period -that ``period`` contains, ``delete(None, branch_name)`` removes all of the -branch's values, and eternal storage keeps one ``ETERNITY`` value per branch. -After every step both storages hold exactly the model's (branch, period) -keys, with the values last put there. +``put`` stores a value unless no key can hold its branch name and period (a +branch name containing ``:``, or a period whose string form names a period +starting on another day: a year or month anchored mid-month), and then both +storages reject it, with the same error; ``delete(period, branch_name)`` +removes the branch's values for every period that ``period`` contains, +``delete(None, branch_name)`` removes all of the branch's values, and eternal +storage keeps one ``ETERNITY`` value per branch. After every step both +storages hold exactly the model's (branch, period) keys, with the values last +put there. + +A second property restores a disk storage's directory, holding the files of +random keys and of names that are not keys, and checks that ``restore`` reads +back exactly the keys, warns about the other names, and that deletes then +agree with memory. Branch names include ones that contain ``_`` or start with another branch's name followed by ``_``, since disk keys join branch and period with ``_``. @@ -17,6 +26,7 @@ import os import tempfile +import warnings import numpy as np import pytest @@ -37,6 +47,8 @@ "y", "y_2019", "trailing_", + # No key can hold a branch name containing ``:``. + "my:reform", ) @@ -54,12 +66,42 @@ def period_strategy(draw): else: month = draw(st.integers(1, 12)) size = draw(st.integers(1, 3)) - # In-memory storage rejects years and months anchored mid-month, whose - # string form drops the day. - day = draw(st.integers(1, 28)) if unit == periods.DAY else 1 + # Both storages reject years and months anchored mid-month, whose string + # form drops the day (before writing a file, so on Windows too). + if unit == periods.DAY: + day = draw(st.integers(1, 28)) + else: + day = draw(st.one_of(st.just(1), st.integers(2, 28))) return periods.Period((unit, periods.Instant((year, month, day)), size)) +def storable(branch_name, period): + """The model of which keys a storage holds: those whose branch name has no + ``:`` and whose period's string form names a period starting the same + day.""" + return ( + ":" not in branch_name + and periods.period(str(period)).start == periods.period(period).start + ) + + +def put_both(memory, disk, value, period, branch_name): + """Put ``value`` in both storages; whether both stored it. Both store it, + or both raise the same ``ValueError``, as the model says.""" + errors = [] + for storage in (memory, disk): + try: + storage.put(np.array([float(value)]), period, branch_name) + errors.append(None) + except ValueError as error: + errors.append(str(error)) + assert errors[0] == errors[1] + if memory.is_eternal: + period = periods.period(periods.ETERNITY) + assert (errors[0] is None) == storable(branch_name, period) + return errors[0] is None + + put_step = st.tuples( st.just("put"), period_strategy(), @@ -102,9 +144,8 @@ def test_disk_and_memory_storage_delete_the_same_keys(is_eternal, steps): if period is not None and is_eternal: period = periods.period(periods.ETERNITY) if action == "put": - memory.put(np.array([float(value)]), period, branch_name) - disk.put(np.array([float(value)]), period, branch_name) - model[(branch_name, str(period))] = float(value) + if put_both(memory, disk, value, period, branch_name): + model[(branch_name, str(period))] = float(value) else: memory.delete(period, branch_name) disk.delete(period, branch_name) @@ -119,3 +160,67 @@ def test_disk_and_memory_storage_delete_the_same_keys(is_eternal, steps): } assert stored(disk) == stored(memory) == model + + +# Names of ``.npy`` files that are not keys: no ``_``, no period after the last +# ``_``, or a period not in its string form. (Not ``eternity``: on a +# case-insensitive file system that is the ``ETERNITY`` key's file.) +NON_CANONICAL_PERIODS = ["2025-3", "2025-03-1", "x", ""] +if os.name != "nt": + NON_CANONICAL_PERIODS += ["year:2025", "month:2025-01:12", "day:2025-03-01"] +not_key_names = st.one_of( + st.text(alphabet="abcxyz0123456789-", min_size=1, max_size=8), + st.builds( + "{}_{}".format, + st.sampled_from([name for name in BRANCH_NAMES if ":" not in name]), + st.one_of( + st.sampled_from(NON_CANONICAL_PERIODS), + st.text(alphabet="abcxyz", min_size=1, max_size=6), + ), + ), +).map(lambda name: name + ".npy") + + +@hypothesis.settings( + max_examples=100, + deadline=None, + suppress_health_check=[hypothesis.HealthCheck.too_slow], +) +@hypothesis.given( + is_eternal=st.booleans(), + puts=st.lists(put_step, max_size=12), + not_keys=st.lists(not_key_names, max_size=4, unique=True), + deletes=st.lists(delete_step, max_size=4), +) +def test_restore_reads_back_the_keys_and_only_the_keys( + is_eternal, puts, not_keys, deletes +): + memory = InMemoryStorage(is_eternal=is_eternal) + with tempfile.TemporaryDirectory() as directory: + disk = OnDiskStorage( + directory, is_eternal=is_eternal, preserve_storage_dir=True + ) + for _, period, branch_name, value in puts: + put_both(memory, disk, value, period, branch_name) + for name in not_keys: + np.save(os.path.join(directory, name), np.array([-1.0])) + + restored = OnDiskStorage( + directory, is_eternal=is_eternal, preserve_storage_dir=True + ) + with warnings.catch_warnings(record=True) as caught: + warnings.simplefilter("always") + restored.restore() + + assert stored(restored) == stored(memory) + if not_keys: + (warning,) = caught + for name in not_keys: + assert name in str(warning.message) + else: + assert caught == [] + + for _, period, branch_name, _ in deletes: + memory.delete(period, branch_name) + restored.delete(period, branch_name) + assert stored(restored) == stored(memory) diff --git a/tests/core/test_known_periods_branch_visibility.py b/tests/core/test_known_periods_branch_visibility.py index df7ea35c1..9bdfd1343 100644 --- a/tests/core/test_known_periods_branch_visibility.py +++ b/tests/core/test_known_periods_branch_visibility.py @@ -490,3 +490,18 @@ def test_restored_branch_dump_calculates_what_the_branch_calculates(system): assert only(restored.calculate("uprated_income", 2021)) == pytest.approx( 3_000 * growth(2019, 2021) ) + + +def test_restore_simulation_leaves_out_files_that_are_not_values(system, tmp_path): + simulation = new_simulation(system) + simulation.set_input("uprated_income", 2016, [1_000.0]) + directory = str(tmp_path / "dump") + simulation_dumper.dump_simulation(simulation, directory) + np.save(os.path.join(directory, "uprated_income", "stray.npy"), np.array([9.0])) + + with pytest.warns(UserWarning, match="stray.npy"): + restored = simulation_dumper.restore_simulation(directory, system) + + holder = restored.get_holder("uprated_income") + assert holder.get_known_periods() == [periods.period(2016)] + assert only(holder.get_array(2016)) == 1_000 From 0432ced3067eafa12f02e4b13b3cab1e5186c367 Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Wed, 7 Oct 2026 17:06:53 -0400 Subject: [PATCH 07/15] Reject unparseable storage periods before writing values Cover the remaining #565 review witness, periods.period(999), on both string-key storage backends. Preserve supported twelve-month aliases and existing derived values after rejected writes. Extend the differential property across the early-year boundary and twelve-unit periods. Co-Authored-By: Claude Opus 5.5 --- .../fix-uprating-branch-visibility.fixed.md | 2 +- docs/python_api/data_storage.md | 7 +++ .../data_storage/on_disk_storage.py | 13 +++--- .../data_storage/storage_keys.py | 26 +++++++---- ...t_disk_storage_delete_contained_periods.py | 44 +++++++++++++++++++ .../core/test_disk_storage_delete_property.py | 15 ++++--- 6 files changed, 84 insertions(+), 23 deletions(-) diff --git a/changelog.d/fix-uprating-branch-visibility.fixed.md b/changelog.d/fix-uprating-branch-visibility.fixed.md index 708932a71..268404cc9 100644 --- a/changelog.d/fix-uprating-branch-visibility.fixed.md +++ b/changelog.d/fix-uprating-branch-visibility.fixed.md @@ -1 +1 @@ -`Holder.get_known_periods(branch_name)` lists only the periods that branch reads, and `Simulation` uses it, so values stored only under other branches no longer change what a branch calculates or caches; on-disk storage now parses branch names containing `_`, deletes the periods within a deleted period as in-memory storage does, rejects the keys in-memory storage rejects (years and months anchored mid-month, branch names containing `:`) and leaves files that are not its keys out of `restore`; and `dump_simulation` saves the values the dumped branch reads, with their derived marks. +Scope known periods and dumps to the branch's readable values while preserving derived marks, parse underscore-containing disk branch names safely, delete contained disk periods consistently with memory, reject ambiguous or unparseable storage keys before writing, and warn and skip malformed restored filenames. diff --git a/docs/python_api/data_storage.md b/docs/python_api/data_storage.md index 2227f4999..20b87a1e7 100644 --- a/docs/python_api/data_storage.md +++ b/docs/python_api/data_storage.md @@ -22,6 +22,13 @@ The `policyengine_core.data_storage` module contains the classes that are used t ## Storage keys +Both storage backends reject branch names containing `:` and month or year +periods anchored after the first day of a month. They also reject periods whose +string form cannot be parsed back, such as the year `0999`, before storing a +value. Existing aliases, including twelve-month periods serialized as one year, +remain supported. Disk restoration warns and skips filenames without a branch +separator and a canonical period string. + ```{eval-rst} .. automodule:: policyengine_core.data_storage.storage_keys :members: diff --git a/policyengine_core/data_storage/on_disk_storage.py b/policyengine_core/data_storage/on_disk_storage.py index 1dcdfc328..bfbfafd5a 100644 --- a/policyengine_core/data_storage/on_disk_storage.py +++ b/policyengine_core/data_storage/on_disk_storage.py @@ -62,9 +62,11 @@ def _split_key(key: str) -> tuple: def _is_key(key: str) -> bool: - """Whether ``key`` is one ``put`` writes: ``f"{branch_name}_{period}"``, - with the period in its string form, so that it splits back into the - branch name and the period it was written for.""" + """Whether ``key`` ends with a canonical period after the branch separator. + + Legacy branch names and eternal-storage files retain their restore + behavior; this check only validates the period's spelling. + """ _, separator, period = key.rpartition("_") if not separator: return False @@ -408,9 +410,8 @@ def restore(self) -> None: """Read back the values stored in this storage's directory: for each key, the file named for it (not those in ``REPLACEMENTS_DIR``). - A ``.npy`` file whose name is not a key ``put`` writes (a branch - name, ``_`` and a period in its string form) holds no value of this - storage: it is left out, with a warning. + A ``.npy`` file whose name lacks a branch separator and a canonical + period string is left out, with a warning. This storage writes over those files, as over those it wrote, until it is cloned or copied, or this process forks: from then on it writes diff --git a/policyengine_core/data_storage/storage_keys.py b/policyengine_core/data_storage/storage_keys.py index bb5575026..d17401b88 100644 --- a/policyengine_core/data_storage/storage_keys.py +++ b/policyengine_core/data_storage/storage_keys.py @@ -2,10 +2,9 @@ Both storages key a value by its branch name and its period's string form: ``f"{branch_name}:{period}"`` in memory, ``f"{branch_name}_{period}"`` on disk -(the name of the value's file). ``check_storable`` rejects what such a key -cannot hold, the same way for both storages, so whether a holder accepts a -value never depends on which of the two it goes to (``Holder._set`` chooses by -how much memory is in use). +(the name of the value's file). ``check_storable`` enforces the same branch +and period serialization restrictions in both storages (``Holder._set`` +chooses between them by how much memory is in use). """ from policyengine_core import periods @@ -13,14 +12,16 @@ def check_storable(branch_name: str, period: Period) -> None: - """Raise ``ValueError`` if no storage key can hold ``branch_name`` and - ``period``. + """Raise ``ValueError`` for an ambiguous branch or period serialization. - Two cases: a branch name containing ``:``, which separates the parts of an + Reject a branch name containing ``:``, which separates the parts of an in-memory key, and a year or month anchored mid-month, whose string form drops the day (``month:2025-03-15`` is ``2025-03``), so that its key would - be the key of another period. See policyengine-core#526 for the - structured-key refactor that will lift these restrictions. + be the key of another period. Reject unparseable period strings too, + including years below 1000 whose string form loses leading zeroes. + Supported aliases, such as twelve months serialized as one year, retain + their existing meaning. See policyengine-core#526 for the structured-key + refactor. """ if ":" in branch_name: raise ValueError( @@ -33,3 +34,10 @@ def check_storable(branch_name: str, period: Period) -> None: f"{period.start}): its string form is lossy (see " "policyengine-core#526)." ) + try: + periods.period(str(period)) + except ValueError as error: + raise ValueError( + f"Cannot cache period {period!r}: its string form cannot be parsed " + "(see policyengine-core#526)." + ) from error diff --git a/tests/core/test_disk_storage_delete_contained_periods.py b/tests/core/test_disk_storage_delete_contained_periods.py index cad439f5f..0505f0f99 100644 --- a/tests/core/test_disk_storage_delete_contained_periods.py +++ b/tests/core/test_disk_storage_delete_contained_periods.py @@ -376,6 +376,49 @@ def test_mid_month_input_is_rejected_whichever_storage_it_goes_to(on_disk): ] +@pytest.mark.parametrize("backend", BACKENDS) +@pytest.mark.parametrize("period", [999, "0999", "0999-03", "0999-03-15"]) +def test_put_rejects_unparseable_period_without_changing_values( + make_storage, backend, period +): + # The full #565 review also found put(periods.period(999)): its "999" + # key cannot be parsed by containment deletion or restored from disk. + storage = make_storage(backend) + storage.put(np.array([500.0]), "2025-03", "no_salt", derived=True) + period = periods.period(period) + + with pytest.raises(ValueError, match="string form cannot be parsed"): + storage.put(np.array([9.0]), period, "no_salt") + + assert storage.get(period, "no_salt") is None + assert known(storage) == [("no_salt", "2025-03")] + assert storage.is_derived("2025-03", "no_salt") + np.testing.assert_array_equal(storage.get("2025-03", "no_salt"), [500.0]) + storage.delete(periods.period(2025), "no_salt") + assert known(storage) == [] + + +@pytest.mark.parametrize("backend", BACKENDS) +@pytest.mark.parametrize("month", [1, 3]) +def test_twelve_month_alias_remains_storable(make_storage, backend, month): + period = periods.Period((periods.MONTH, periods.Instant((2025, month, 1)), 12)) + if os.name == "nt" and ":" in str(period): + pytest.skip("Windows file names cannot contain ':'") + storage = make_storage(backend) + + storage.put(np.array([500.0]), period, "no_salt", derived=True) + + alias = periods.period(str(period)) + assert alias.unit == periods.YEAR + assert alias.start == period.start + assert alias.stop == period.stop + assert storage.get_known_branch_periods() == [("no_salt", alias)] + assert storage.is_derived(alias, "no_salt") + np.testing.assert_array_equal(storage.get(alias, "no_salt"), [500.0]) + storage.delete(period, "no_salt") + assert known(storage) == [] + + # ----- Files restored from a directory (#565 review, finding 2) --------------- # ``restore`` read back every ``.npy`` file in the directory. One whose name @@ -390,6 +433,7 @@ def test_mid_month_input_is_rejected_whichever_storage_it_goes_to(on_disk): "_.npy", "default_2025-3.npy", "default_2025-03-1.npy", + "default_999.npy", ] if os.name != "nt": NOT_KEYS += ["default_month:2025-01:12.npy", "no_salt_year:2025:1.npy"] diff --git a/tests/core/test_disk_storage_delete_property.py b/tests/core/test_disk_storage_delete_property.py index f384ec948..8f92bd761 100644 --- a/tests/core/test_disk_storage_delete_property.py +++ b/tests/core/test_disk_storage_delete_property.py @@ -57,7 +57,7 @@ def period_strategy(draw): if draw(st.integers(0, 19)) == 0: return periods.period(periods.ETERNITY) unit = draw(st.sampled_from([periods.YEAR, periods.MONTH, periods.DAY])) - year = draw(st.integers(2023, 2026)) + year = draw(st.one_of(st.sampled_from([999, 1000]), st.integers(2023, 2026))) if os.name == "nt": # Windows file names cannot contain the ``:`` of a multi-unit or # rolling-year period's string form. @@ -65,7 +65,7 @@ def period_strategy(draw): size = 1 else: month = draw(st.integers(1, 12)) - size = draw(st.integers(1, 3)) + size = draw(st.sampled_from([1, 2, 3, 12])) # Both storages reject years and months anchored mid-month, whose string # form drops the day (before writing a file, so on Windows too). if unit == periods.DAY: @@ -79,10 +79,11 @@ def storable(branch_name, period): """The model of which keys a storage holds: those whose branch name has no ``:`` and whose period's string form names a period starting the same day.""" - return ( - ":" not in branch_name - and periods.period(str(period)).start == periods.period(period).start - ) + try: + parsed = periods.period(str(period)) + except ValueError: + return False + return ":" not in branch_name and parsed.start == periods.period(period).start def put_both(memory, disk, value, period, branch_name): @@ -165,7 +166,7 @@ def test_disk_and_memory_storage_delete_the_same_keys(is_eternal, steps): # Names of ``.npy`` files that are not keys: no ``_``, no period after the last # ``_``, or a period not in its string form. (Not ``eternity``: on a # case-insensitive file system that is the ``ETERNITY`` key's file.) -NON_CANONICAL_PERIODS = ["2025-3", "2025-03-1", "x", ""] +NON_CANONICAL_PERIODS = ["2025-3", "2025-03-1", "999", "x", ""] if os.name != "nt": NON_CANONICAL_PERIODS += ["year:2025", "month:2025-01:12", "day:2025-03-01"] not_key_names = st.one_of( From 92f7fd98c1d01cd5e2fe06b56bf8866a3056a2b8 Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Mon, 19 Oct 2026 02:55:40 -0400 Subject: [PATCH 08/15] Reconcile master storage tests with branch validation and deletion Keep the PR's rejection of unparseable early-year periods and malformed restored keys. Update master expectations for contained-period disk deletion while retaining input bookkeeping and clone isolation coverage. Production behavior is unchanged. --- tests/core/test_branch_helper_invalidation.py | 11 ++- tests/core/test_disk_period_filenames.py | 24 ++++-- tests/core/test_registered_input_deletion.py | 34 +++++++- tests/core/test_set_input_helper_order.py | 29 +++---- tests/core/test_user_input_keys_soundness.py | 79 ++++++++++++------- 5 files changed, 115 insertions(+), 62 deletions(-) diff --git a/tests/core/test_branch_helper_invalidation.py b/tests/core/test_branch_helper_invalidation.py index ba66b7847..deca3f205 100644 --- a/tests/core/test_branch_helper_invalidation.py +++ b/tests/core/test_branch_helper_invalidation.py @@ -78,15 +78,14 @@ def build(): january = year.first_month months = [january.offset(index) for index in range(12)] parent = build() + if year_string == "0025": + with pytest.raises(ValueError, match="string form cannot be parsed"): + parent.calculate("helper_independent", year) + assert parent.get_holder("helper_independent").get_known_periods() == [] + return parent.calculate("helper_independent", year) parent.set_input("helper_flow", january, [12.0, 24.0]) parent_months = [january] - if year_string == "0025": - # Early-year fallback getters have a separate legacy parser issue. - # Supply every month so this case exercises history and input replay. - for month in months[1:]: - parent.set_input("helper_flow", month, [0.0, 0.0]) - parent_months = months parent_inputs = {("helper_flow", "default", month) for month in parent_months} np.testing.assert_array_equal(parent.calculate("helper_flow", year), [12, 24]) np.testing.assert_array_equal(parent.calculate("helper_result", year), [24, 48]) diff --git a/tests/core/test_disk_period_filenames.py b/tests/core/test_disk_period_filenames.py index 5c68eeb5a..5f9407a77 100644 --- a/tests/core/test_disk_period_filenames.py +++ b/tests/core/test_disk_period_filenames.py @@ -38,7 +38,7 @@ def checked_save(file, *args, **kwargs): monkeypatch.setattr(np, "save", checked_save) -@pytest.mark.parametrize("period", COMPOUND_PERIODS + EARLY_DAYS, ids=str) +@pytest.mark.parametrize("period", COMPOUND_PERIODS + EARLY_DAYS[-1:], ids=str) def test_compound_period_has_a_portable_disk_filename( tmp_path, portable_writes, period ): @@ -54,6 +54,19 @@ def test_compound_period_has_a_portable_disk_filename( assert set(restored._files) == {f"default_{period}"} +@pytest.mark.parametrize("period", EARLY_DAYS[:-1], ids=str) +def test_unparseable_early_period_is_rejected_before_writing( + tmp_path, portable_writes, period +): + storage = OnDiskStorage(str(tmp_path), preserve_storage_dir=True) + + with pytest.raises(ValueError, match="string form cannot be parsed"): + storage.put(np.array([42.0]), period) + + assert not storage._files + assert not list(tmp_path.iterdir()) + + @pytest.mark.parametrize("period", COMPOUND_PERIODS, ids=str) def test_compound_period_clones_use_portable_replacement_paths( tmp_path, portable_writes, period @@ -130,14 +143,15 @@ def test_restore_still_reads_legacy_compound_period_filename(tmp_path): @pytest.mark.parametrize("name", ["notes", "notes;percent%3A"]) -def test_restore_tolerates_basenames_without_a_key_separator(tmp_path, name): +def test_restore_skips_basenames_without_a_key_separator(tmp_path, name): np.save(tmp_path / f"{name}.npy", np.array([42.0])) restored = OnDiskStorage(str(tmp_path), preserve_storage_dir=True) - restored.restore() + with pytest.warns(UserWarning, match="Not restoring"): + restored.restore() - assert set(restored._files) == {name} - np.testing.assert_array_equal(restored._decode_file(restored._files[name]), [42.0]) + assert not restored._files + np.testing.assert_array_equal(np.load(tmp_path / f"{name}.npy"), [42.0]) @pytest.mark.skipif( diff --git a/tests/core/test_registered_input_deletion.py b/tests/core/test_registered_input_deletion.py index 0d1f02ab1..1fb9db1a9 100644 --- a/tests/core/test_registered_input_deletion.py +++ b/tests/core/test_registered_input_deletion.py @@ -1,4 +1,4 @@ -"""Restored input records follow deletion even for early-year storage keys.""" +"""Restored input records respect storage period validation and deletion.""" import numpy as np import pytest @@ -34,6 +34,15 @@ def test_deleted_registered_early_year_input_is_forgotten(year, variable_type, o (variable_type.definition_period, periods.instant((year, 1, 1)), 1) ) storage = holder._disk_storage if on_disk else holder._memory_storage + if year < 1000: + input_keys = simulation._user_input_keys.copy() + with pytest.raises(ValueError, match="string form cannot be parsed"): + storage.put(np.array([42.0]), period) + assert not storage.has(period) + assert not storage._sequence_numbers + assert simulation._user_input_keys == input_keys + assert not hasattr(holder, "_user_input_storage") + return storage.put(np.array([42.0]), period) key = (variable_type.__name__, "default", period) simulation._user_input_keys.add(key) @@ -65,9 +74,10 @@ def test_registered_deletion_matches_set_input_for_period_aliases( ): """Register-only and set_input paths forget the same normalized slot. - Deletion removes the input from the record regardless of the stored - period's unit, size, early year, alias, tier, or underscored branch. - A write afterward remains a cache value and is cleared by a reform. + Unsupported early years are rejected before registration. For supported + years, deletion removes the input regardless of the stored period's unit, + size, alias, tier, or underscored branch. A write afterward remains a + cache value and is cleared by a reform. """ restored = _raw_simulation(variable_type, on_disk) tracked = _raw_simulation(variable_type, on_disk) @@ -75,6 +85,22 @@ def test_registered_deletion_matches_set_input_for_period_aliases( branch = "restore_branch" holder = restored.get_holder(variable_type.__name__) storage = holder._disk_storage if on_disk else holder._memory_storage + if year < 1000: + simulations = (restored, tracked) + input_keys = [simulation._user_input_keys.copy() for simulation in simulations] + with pytest.raises(ValueError, match="string form cannot be parsed"): + storage.put(np.array([42.0]), period, branch) + with pytest.raises(ValueError, match="string form cannot be parsed"): + tracked.get_holder(variable_type.__name__).set_input(period, [42.0], branch) + for simulation, before in zip(simulations, input_keys): + current = simulation.get_holder(variable_type.__name__) + tier = current._disk_storage if on_disk else current._memory_storage + assert not tier.has(period, branch) + assert not tier._sequence_numbers + assert simulation._user_input_keys == before + assert simulation._user_input_contexts == [] + assert not hasattr(current, "_user_input_storage") + return storage.put(np.array([42.0]), period, branch) stored_period = holder._storage_period(period) key = (variable_type.__name__, branch, stored_period) diff --git a/tests/core/test_set_input_helper_order.py b/tests/core/test_set_input_helper_order.py index b4fb8fc3b..4a1620de9 100644 --- a/tests/core/test_set_input_helper_order.py +++ b/tests/core/test_set_input_helper_order.py @@ -581,29 +581,24 @@ def test_repeated_helper_clears_cache_shadow_when_every_month_has_a_disk_input( @pytest.mark.parametrize("on_disk", [False, True]) -def test_helper_drops_overlapping_early_year_cache(on_disk): +def test_helper_rejects_early_year_input_and_cache_without_recording(on_disk): simulation = build_simulation(on_disk=on_disk) holder = simulation.get_holder("flow_m") year = periods.period("0025") - holder.put_in_cache(np.array([5.0, 6.0]), year, derived=True) - # Plain reads do not fast-cache an aggregate today; model an existing - # fast-cache entry to prove cleanup uses the early-year Period itself. - simulation._fast_cache[("flow_m", year)] = np.array([5.0, 6.0]) - fresh = build_simulation() + input_keys = simulation._user_input_keys.copy() - set_input_divide_by_period(holder, year, YEARLY_INPUT) - fresh.set_input("flow_m", year, YEARLY_INPUT) + with pytest.raises(ValueError, match="string form cannot be parsed"): + holder.put_in_cache(np.array([5.0, 6.0]), year, derived=True) + with pytest.raises(ValueError, match="string form cannot be parsed"): + set_input_divide_by_period(holder, year, YEARLY_INPUT) - assert holder.get_array(year) is None + assert holder.get_known_periods() == [] + storage = holder._disk_storage if on_disk else holder._memory_storage + assert not storage._sequence_numbers assert ("flow_m", year) not in simulation._fast_cache - january = year.start.period(periods.MONTH) - for month in [january.offset(index) for index in range(12)]: - np.testing.assert_array_equal( - simulation.calculate("flow_m", month), fresh.calculate("flow_m", month) - ) - np.testing.assert_array_equal( - simulation.calculate("flow_m", year), fresh.calculate("flow_m", year) - ) + assert simulation._user_input_keys == input_keys + assert not hasattr(holder, "_user_input_storage") + assert not getattr(simulation, "_user_input_contexts", []) def test_every_stored_value_counts_as_an_input_without_a_record(): diff --git a/tests/core/test_user_input_keys_soundness.py b/tests/core/test_user_input_keys_soundness.py index 9fad83d1e..f0d769233 100644 --- a/tests/core/test_user_input_keys_soundness.py +++ b/tests/core/test_user_input_keys_soundness.py @@ -174,7 +174,7 @@ def test_memory_input_wins_over_a_conflicting_disk_cache(invalidate): @pytest.mark.parametrize("variable_type", [raw_year, raw_month, raw_day]) @pytest.mark.parametrize("entrypoint", ["simulation", "holder"]) @pytest.mark.parametrize("on_disk", [False, True]) -def test_early_year_input_is_stored_and_recorded( +def test_early_year_input_respects_storage_period_validation( year, variable_type, entrypoint, on_disk ): system = CountryTaxBenefitSystem() @@ -191,6 +191,21 @@ def test_early_year_input_is_stored_and_recorded( else: period = periods.Period((periods.DAY, periods.instant((year, 3, 17)), 2)) + if year < 1000: + input_keys = simulation._user_input_keys.copy() + with pytest.raises(ValueError, match="string form cannot be parsed"): + if entrypoint == "simulation": + simulation.set_input(variable, period, [120.0]) + else: + holder.set_input(period, [120.0]) + storage = holder._disk_storage if on_disk else holder._memory_storage + assert not storage.has(period) + assert not storage._sequence_numbers + assert simulation._user_input_keys == input_keys + assert not hasattr(holder, "_user_input_storage") + assert simulation._user_input_contexts == [] + return + if entrypoint == "simulation": simulation.set_input(variable, period, [120.0]) else: @@ -210,19 +225,23 @@ def test_early_year_input_is_stored_and_recorded( @pytest.mark.parametrize("on_disk", [False, True]) -def test_zero_padded_early_year_string_records_the_stored_period(on_disk): +def test_zero_padded_early_year_string_is_rejected_before_recording(on_disk): system = CountryTaxBenefitSystem() system.add_variable(raw_year) simulation = _simulation(system) if on_disk: _store_on_disk(simulation, "raw_year") - simulation.set_input("raw_year", "0999", [123.0]) + input_keys = simulation._user_input_keys.copy() + holder = simulation.get_holder("raw_year") + with pytest.raises(ValueError, match="string form cannot be parsed"): + simulation.set_input("raw_year", "0999", [123.0]) - assert _key("raw_year", periods.period(999)) in simulation._user_input_keys - np.testing.assert_array_equal( - simulation.get_holder("raw_year").get_array(periods.period(999)), [123.0] - ) + storage = holder._disk_storage if on_disk else holder._memory_storage + assert not storage.has(periods.period(999)) + assert not storage._sequence_numbers + assert simulation._user_input_keys == input_keys + assert not hasattr(holder, "_user_input_storage") assert simulation._user_input_contexts == [] @@ -250,7 +269,7 @@ def test_deleted_memory_input_does_not_turn_a_disk_calculation_into_an_input(): assert _key("rent", JANUARY) not in simulation._user_input_keys -def test_deleted_memory_input_keeps_a_surviving_real_disk_input(): +def test_contained_deletion_forgets_both_real_input_tiers(): simulation = _simulation() _store_on_disk(simulation, "rent") holder = simulation.get_holder("rent") @@ -263,15 +282,16 @@ def test_deleted_memory_input_keeps_a_surviving_real_disk_input(): simulation.delete_arrays("rent", "2025") assert holder._memory_storage.get(JANUARY) is None - assert _key("rent", JANUARY) in simulation._user_input_keys - assert simulation.to_input_dataframe()["rent__2025-01"].tolist() == [500.0] + assert holder._disk_storage.get(JANUARY) is None + assert _key("rent", JANUARY) not in simulation._user_input_keys + assert "rent__2025-01" not in simulation.to_input_dataframe().columns simulation.apply_reform(NoopReform) assert holder._memory_storage.get(JANUARY) is None - np.testing.assert_array_equal(holder._disk_storage.get(JANUARY), [500.0]) - assert _key("rent", JANUARY) in simulation._user_input_keys - assert simulation.to_input_dataframe()["rent__2025-01"].tolist() == [500.0] + assert holder._disk_storage.get(JANUARY) is None + assert _key("rent", JANUARY) not in simulation._user_input_keys + assert "rent__2025-01" not in simulation.to_input_dataframe().columns @pytest.mark.parametrize("on_disk", [False, True]) @@ -296,7 +316,7 @@ def test_reform_does_not_replay_an_unrecorded_carry_over_input(on_disk): ) -def test_registered_disk_input_survives_a_new_memory_input_and_clone_deletion(): +def test_registered_disk_input_is_forgotten_on_contained_clone_deletion(): """Restoration can register an input without holder tier metadata (#576).""" simulation = _simulation() _store_on_disk(simulation, "rent") @@ -310,18 +330,19 @@ def test_registered_disk_input_survives_a_new_memory_input_and_clone_deletion(): clone = simulation.clone() clone.delete_arrays("rent", "2025") - assert _key("rent", JANUARY) in clone._user_input_keys - assert clone.to_input_dataframe()["rent__2025-01"].tolist() == [500.0] + assert _key("rent", JANUARY) not in clone._user_input_keys + assert "rent__2025-01" not in clone.to_input_dataframe().columns + assert clone.get_holder("rent")._disk_storage.get(JANUARY) is None + np.testing.assert_array_equal(holder._memory_storage.get(JANUARY), [700.0]) + np.testing.assert_array_equal(holder._disk_storage.get(JANUARY), [500.0]) assert simulation.to_input_dataframe()["rent__2025-01"].tolist() == [700.0] simulation.delete_arrays("rent", "2025") for sim in (simulation, clone): - assert _key("rent", JANUARY) in sim._user_input_keys - assert sim.to_input_dataframe()["rent__2025-01"].tolist() == [500.0] + assert _key("rent", JANUARY) not in sim._user_input_keys + assert "rent__2025-01" not in sim.to_input_dataframe().columns sim.apply_reform(NoopReform) assert sim.get_holder("rent")._memory_storage.get(JANUARY) is None - np.testing.assert_array_equal( - sim.get_holder("rent")._disk_storage.get(JANUARY), [500.0] - ) + assert sim.get_holder("rent")._disk_storage.get(JANUARY) is None def test_positional_derived_argument_keeps_masters_meaning(): @@ -347,22 +368,22 @@ def test_positional_derived_argument_keeps_masters_meaning(): @hypothesis.settings(max_examples=30, deadline=None) @hypothesis.example( - disk_is_input=False, delete_memory=True, disk_value=0, memory_value=500 + disk_is_input=False, delete_period=True, disk_value=0, memory_value=500 ) @hypothesis.example( - disk_is_input=True, delete_memory=True, disk_value=500, memory_value=700 + disk_is_input=True, delete_period=True, disk_value=500, memory_value=700 ) @hypothesis.example( - disk_is_input=False, delete_memory=False, disk_value=99, memory_value=42 + disk_is_input=False, delete_period=False, disk_value=99, memory_value=42 ) @hypothesis.given( disk_is_input=st.booleans(), - delete_memory=st.booleans(), + delete_period=st.booleans(), disk_value=st.integers(min_value=0, max_value=1_000), memory_value=st.integers(min_value=0, max_value=1_000), ) def test_tier_provenance_replays_the_same_inputs_as_a_fresh_simulation( - disk_is_input, delete_memory, disk_value, memory_value + disk_is_input, delete_period, disk_value, memory_value ): simulation = _simulation() _store_on_disk(simulation, "rent") @@ -373,12 +394,10 @@ def test_tier_provenance_replays_the_same_inputs_as_a_fresh_simulation( holder.put_in_cache(np.array([float(disk_value)]), JANUARY) simulation.memory_config.max_memory_occupation_pc = 101 simulation.set_input("rent", JANUARY, [memory_value]) - if delete_memory: + if delete_period: simulation.delete_arrays("rent", "2025") - expected = ( - (disk_value if disk_is_input else None) if delete_memory else memory_value - ) + expected = None if delete_period else memory_value reference = _simulation() if expected is not None: reference.set_input("rent", JANUARY, [expected]) From 88dd47c7c41d2b45e2f439d3d782b0f0d5ed3528 Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Tue, 20 Oct 2026 19:29:42 -0400 Subject: [PATCH 09/15] Leave root carry-over and uprating defaults uncached --- .../fix-uprating-branch-visibility.fixed.md | 2 +- docs/usage/branches.md | 5 +- policyengine_core/simulations/simulation.py | 23 +++++-- tests/core/test_carry_over_order.py | 38 +++++++---- .../test_known_periods_branch_visibility.py | 66 +++++++++++++++++++ 5 files changed, 114 insertions(+), 20 deletions(-) diff --git a/changelog.d/fix-uprating-branch-visibility.fixed.md b/changelog.d/fix-uprating-branch-visibility.fixed.md index 268404cc9..f1ef6e416 100644 --- a/changelog.d/fix-uprating-branch-visibility.fixed.md +++ b/changelog.d/fix-uprating-branch-visibility.fixed.md @@ -1 +1 @@ -Scope known periods and dumps to the branch's readable values while preserving derived marks, parse underscore-containing disk branch names safely, delete contained disk periods consistently with memory, reject ambiguous or unparseable storage keys before writing, and warn and skip malformed restored filenames. +Scope known periods and dumps to the branch's readable values while preserving derived marks, leave root uprating and carry-over fallback defaults uncached so later inputs take effect, parse underscore-containing disk branch names safely, delete contained disk periods consistently with memory, reject ambiguous or unparseable storage keys before writing, and warn and skip malformed restored filenames. diff --git a/docs/usage/branches.md b/docs/usage/branches.md index 8a103f6fc..0629a27ab 100644 --- a/docs/usage/branches.md +++ b/docs/usage/branches.md @@ -237,7 +237,10 @@ depend on the input is calculated again) but not less, within these limits: input. Input helpers still remove overlapping calculated values of the variable they write. A formula's input for its own period is still returned as described above, and results calculated while an input is being set do not - enter the fast cache. + enter the fast cache. When uprating or automatic carry-over finds no source + input, the root returns the fallback default without caching it. Setting an + earlier input can therefore change the next calculation of that period. + Branches can cache these defaults because later inputs invalidate them. - **Existing child branches keep their values.** An input set on a branch does not reach branches already created from it. - **Values kept elsewhere stay.** A value one simulation calculated from diff --git a/policyengine_core/simulations/simulation.py b/policyengine_core/simulations/simulation.py index 951b286b4..425e4b762 100644 --- a/policyengine_core/simulations/simulation.py +++ b/policyengine_core/simulations/simulation.py @@ -1791,12 +1791,13 @@ def _calculate(self, variable_name: str, period: Period = None) -> ArrayLike: finally: _formula_simulation.reset(token) - # If no result, use the default value and cache it + # If no result, carry/uprate an input or use the default value. if array is None: - if variable.uprating is not None or ( + cross_period_fallback = variable.uprating is not None or ( self.tax_benefit_system.auto_carry_over_input_variables and variable.calculate_output is None - ): + ) + if cross_period_fallback: # The value is uprated or carried over from another # period, or defaults for lack of one: an input set later # for another period of the variable can change it. @@ -1928,10 +1929,20 @@ def _calculate(self, variable_name: str, period: Period = None) -> ArrayLike: # above skips derived periods, so it would not uprate # from one.) return holder.default_array() - else: - array = holder.default_array() - else: + if array is None: array = holder.default_array() + if ( + cross_period_fallback + and getattr(self, "parent_branch", None) is None + ): + # Root inputs do not invalidate dependent values. + # Leave this fallback out of every cache so a later + # input for an earlier period takes effect. An input + # set by the formula for this period still wins. + stored_input = self._input_set_meanwhile( + holder, period, input_state[2] + ) + return array if stored_input is None else stored_input if variable.defined_for is not None: array = np.where(mask, array, variable.default_value) diff --git a/tests/core/test_carry_over_order.py b/tests/core/test_carry_over_order.py index a144fd78f..df1231027 100644 --- a/tests/core/test_carry_over_order.py +++ b/tests/core/test_carry_over_order.py @@ -154,6 +154,10 @@ def test_an_input_ending_after_the_last_date_drops_overlapping_defaults( holder = built.get_holder(name) requested = periods.period("9999-12-31") holder._do_not_store = cache_only + # Seed existing caches: root carry-over fallbacks no longer store defaults. + default = holder.default_array() + holder.put_in_cache(default, requested, derived=True) + built._fast_cache[(name, requested)] = default np.testing.assert_array_equal(built.calculate(name, requested), [0, 0]) assert (name, requested) in built._fast_cache assert holder._memory_storage.has(requested) == (not cache_only) @@ -259,21 +263,27 @@ def test_input_set_after_a_period_defaulted_still_carries(system): np.testing.assert_array_equal(built.calculate("carried", "2014"), [7, 8]) -def test_defaults_are_cached_as_before(system): - """With nothing stored, or only earlier calculated values, the default is - cached (as master cached the value it carried); before a later stored - period it is not (as master returned it).""" +@pytest.mark.parametrize("on_branch", [False, True], ids=["root", "branch"]) +def test_fallback_defaults_are_cached_only_on_branches(system, on_branch): + """Only branches invalidate these defaults when earlier inputs arrive. + Before a later stored period, neither root nor branch caches them.""" built = simulation(system, {}) + if on_branch: + built = built.get_branch("fallback_defaults") holder = built.get_holder("carried_monthly") for month in range(1, 13): for _ in range(2): built.calculate("carried_monthly", f"2012-{month:02d}") for month in range(1, 13): - assert holder.get_array(periods.period(f"2012-{month:02d}")) is not None + period = periods.period(f"2012-{month:02d}") + assert (holder.get_array(period, built.branch_name) is not None) == on_branch + assert (("carried_monthly", period) in built._fast_cache) == on_branch built = simulation(system, {"carried_monthly": {"2013-06": [1, 2]}}) + if on_branch: + built = built.get_branch("fallback_defaults") holder = built.get_holder("carried_monthly") np.testing.assert_array_equal(built.calculate("carried_monthly", "2012-03"), [0, 0]) - assert holder.get_array(periods.period("2012-03")) is None + assert holder.get_array(periods.period("2012-03"), built.branch_name) is None def test_branch_input_is_not_hidden_by_a_value_the_parent_carried(system): @@ -609,19 +619,21 @@ def test_a_period_stored_only_on_another_branch_does_not_carry(system): def test_marks_are_cleared_with_the_values_apply_reform_wipes(system): - built = simulation(system, {}) + built = simulation(system, {"carried": {"2009": [7, 8]}}) built.memory_config = MemoryConfig(max_memory_occupation=0) holder = built.get_holder("carried") holder._disk_storage = holder.create_disk_storage() holder._on_disk_storable = True for year in range(2010, 2020): built.calculate("carried", str(year)) + assert holder.is_derived(periods.period(year)) built.apply_reform(_noop) - assert not holder._memory_storage._inputs + assert holder.get_known_periods() == [periods.period("2009")] + assert holder.get_input_periods() == [periods.period("2009")] assert not holder._disk_storage._derived # A file written for a period whose value was derived, then read back # by rebuilding the index, is an input. - built = simulation(system, {}) + built = simulation(system, {"carried": {"2011": [7, 8]}}) built.memory_config = MemoryConfig(max_memory_occupation=0) holder = built.get_holder("carried") holder._disk_storage = holder.create_disk_storage() @@ -637,17 +649,19 @@ def test_marks_are_cleared_with_the_values_apply_reform_wipes(system): def test_marks_are_cleared_with_in_memory_values_apply_reform_wipes(system): - built = simulation(system, {}) + built = simulation(system, {"carried": {"2009": [7, 8]}}) holder = built.get_holder("carried") for year in range(2010, 2020): built.calculate("carried", str(year)) + assert holder.is_derived(periods.period(year)) built.apply_reform(_noop) - assert not holder._memory_storage._inputs + assert holder.get_known_periods() == [periods.period("2009")] + assert holder.get_input_periods() == [periods.period("2009")] def test_rebuilding_a_disk_index_reads_files_as_inputs(system): """A derived file replaced by an input, then re-indexed, is an input.""" - built = simulation(system, {}) + built = simulation(system, {"carried": {"2011": [7, 8]}}) built.memory_config = MemoryConfig(max_memory_occupation=0) holder = built.get_holder("carried") holder._disk_storage = holder.create_disk_storage() diff --git a/tests/core/test_known_periods_branch_visibility.py b/tests/core/test_known_periods_branch_visibility.py index 9bdfd1343..2d51e4082 100644 --- a/tests/core/test_known_periods_branch_visibility.py +++ b/tests/core/test_known_periods_branch_visibility.py @@ -35,6 +35,8 @@ import os import tempfile import warnings +from types import SimpleNamespace +from unittest.mock import Mock import numpy as np import pytest @@ -46,6 +48,7 @@ from policyengine_core.experimental import MemoryConfig from policyengine_core.model_api import YEAR, Variable from policyengine_core.parameters import ParameterNode +from policyengine_core.reforms import Reform from policyengine_core.simulations import SimulationBuilder from policyengine_core.tools import simulation_dumper @@ -189,6 +192,69 @@ def test_carry_over_ignores_later_period_on_unrelated_branch(carry_over_system): assert only(simulation.calculate("carried_income", 2020)) == 1_000 +def test_root_default_before_unreadable_future_input_does_not_hide_later_input( + carry_over_system, +): + simulation = new_simulation(carry_over_system) + store(simulation, "carried_income", 2022, 9_000.0, "other") + + assert only(simulation.calculate("carried_income", 2020)) == 0 + + simulation.set_input("carried_income", 2019, [123]) + + assert only(simulation.calculate("carried_income", 2020)) == 123 + + +@pytest.mark.parametrize( + "variable, auto_carry_over", + [("carried_income", True), ("uprated_income", False)], +) +def test_root_fallback_default_is_not_kept_in_any_cache( + variable, auto_carry_over, monkeypatch, tmp_path +): + import policyengine_core.simulations.simulation as simulation_module + + simulation = new_simulation(build_system(auto_carry_over)) + simulation.dataset = SimpleNamespace(file_path=tmp_path / "dataset.h5", name="test") + cache = Mock() + cache.get_cache_path.return_value = tmp_path / "missing-cache.h5" + monkeypatch.setattr(simulation_module, "SimulationMacroCache", lambda system: cache) + monkeypatch.setattr(simulation, "check_macro_cache", lambda *args: True) + + assert only(simulation.calculate(variable, 2020)) == 0 + + period = periods.period(2020) + assert simulation.get_holder(variable).get_array(period) is None + assert (variable, period) not in simulation._fast_cache + cache.set_cache_value.assert_not_called() + + simulation.set_input(variable, 2019, [123]) + expected = 123 * growth(2019, 2020) if variable == "uprated_income" else 123 + assert only(simulation.calculate(variable, 2020)) == pytest.approx(expected) + assert simulation.get_holder(variable).get_array(period) is not None + cache.set_cache_value.assert_called_once() + + +def test_root_uprating_fallback_uses_input_set_during_formula(): + class self_supplied_income(Reform): + def apply(self): + class uprated_income(Variable): + value_type = float + entity = Person + definition_period = YEAR + + def formula(person, period): + person.simulation.set_input("uprated_income", period, [123]) + + self.update_variable(uprated_income) + + simulation = new_simulation(self_supplied_income(build_system())) + + assert only(simulation.calculate("uprated_income", 2020)) == 123 + holder = simulation.get_holder("uprated_income") + assert not holder.is_derived(periods.period(2020)) + + # Invariant: values stored only under branches a simulation cannot read do not # change what it calculates, nor what it stores where it reads. Checked # exhaustively against a simulation that never had those values, for every From a5f2f622362ea2e08dbf52f1c798c2a5dd2f91e5 Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Thu, 22 Oct 2026 03:51:26 -0400 Subject: [PATCH 10/15] Keep converged root fallback results out of caches --- .../fix-uprating-branch-visibility.fixed.md | 2 +- docs/usage/branches.md | 5 +- policyengine_core/simulations/simulation.py | 32 +++-- tests/core/test_branch_input_fixed_point.py | 113 ++++++++++++++++++ .../test_branch_input_fixed_point_property.py | 43 +++++++ 5 files changed, 185 insertions(+), 10 deletions(-) diff --git a/changelog.d/fix-uprating-branch-visibility.fixed.md b/changelog.d/fix-uprating-branch-visibility.fixed.md index f1ef6e416..2abb0cac1 100644 --- a/changelog.d/fix-uprating-branch-visibility.fixed.md +++ b/changelog.d/fix-uprating-branch-visibility.fixed.md @@ -1 +1 @@ -Scope known periods and dumps to the branch's readable values while preserving derived marks, leave root uprating and carry-over fallback defaults uncached so later inputs take effect, parse underscore-containing disk branch names safely, delete contained disk periods consistently with memory, reject ambiguous or unparseable storage keys before writing, and warn and skip malformed restored filenames. +Scope known periods and dumps to the branch's readable values while preserving derived marks, leave root uprating and carry-over fallback defaults uncached even after branch mutations converge so later inputs take effect, parse underscore-containing disk branch names safely, delete contained disk periods consistently with memory, reject ambiguous or unparseable storage keys before writing, and warn and skip malformed restored filenames. diff --git a/docs/usage/branches.md b/docs/usage/branches.md index 0629a27ab..d4dda2a5d 100644 --- a/docs/usage/branches.md +++ b/docs/usage/branches.md @@ -238,8 +238,9 @@ depend on the input is calculated again) but not less, within these limits: variable they write. A formula's input for its own period is still returned as described above, and results calculated while an input is being set do not enter the fast cache. When uprating or automatic carry-over finds no source - input, the root returns the fallback default without caching it. Setting an - earlier input can therefore change the next calculation of that period. + input, the root returns the fallback default without caching it, including + when repeated branch mutations converge to a fixed point. Setting an earlier + input can therefore change the next calculation of that period. Branches can cache these defaults because later inputs invalidate them. - **Existing child branches keep their values.** An input set on a branch does not reach branches already created from it. diff --git a/policyengine_core/simulations/simulation.py b/policyengine_core/simulations/simulation.py index 425e4b762..c9e598944 100644 --- a/policyengine_core/simulations/simulation.py +++ b/policyengine_core/simulations/simulation.py @@ -283,6 +283,7 @@ class _Frame: "stale_parent", "unobserved_activity", "attempt_cache", + "uncached_fallbacks", ) def __init__(self, simulation: "Simulation", stale_parent: bool = False): @@ -307,6 +308,9 @@ def restart(self) -> None: # Restricted children share this dictionary only until this attempt # restarts. Foreign results never enter persistent caches this way. self.attempt_cache = {} + # Cache eligibility belongs to the requested value, not to every + # formula sharing this frame. A retry may find a real source input. + self.uncached_fallbacks = set() # Retain only creations whose original registration still names # the same object. Saved or replaced snapshots keep their identities. self.carried_branches = _registered_creations( @@ -1487,6 +1491,8 @@ def calculate( transition = frame.transition_signature(result) if transition is not None and transition == previous_transition: frame.settle_reads() + if (variable_name, period) in frame.uncached_fallbacks: + break # This run repeated the complete input transition and its # result. Preserve holder input/blacklist rules and the # existing fast-cache guard when keeping the fixed point. @@ -1657,6 +1663,10 @@ def _calculate(self, variable_name: str, period: Period = None) -> ArrayLike: if variable_name not in self.tax_benefit_system.variables: raise ValueError(f"Variable {variable_name} does not exist.") input_state = self._calculation_start() + frames = _calculation_frames.get() + frame = frames[-1] if frames and frames[-1].simulation is self else None + if frame is not None: + frame.uncached_fallbacks.discard((variable_name, period)) # Lets a custom ``set_input`` handler tell whether it calculated. self._calculations_started = self._calculations_started + 1 population = self.get_variable_population(variable_name) @@ -1682,7 +1692,6 @@ def _calculate(self, variable_name: str, period: Period = None) -> ArrayLike: if cached_array is not None: return cached_array - frames = _calculation_frames.get() if frames and frames[-1].simulation is self: cached_array = frames[-1].recall(variable_name, period) if cached_array is not None: @@ -1928,9 +1937,8 @@ def _calculate(self, variable_name: str, period: Period = None) -> ArrayLike: # whether a value is stored sees. (The uprating path # above skips derived periods, so it would not uprate # from one.) - return holder.default_array() + return self._uncached_fallback(holder, period, input_state) if array is None: - array = holder.default_array() if ( cross_period_fallback and getattr(self, "parent_branch", None) is None @@ -1939,10 +1947,8 @@ def _calculate(self, variable_name: str, period: Period = None) -> ArrayLike: # Leave this fallback out of every cache so a later # input for an earlier period takes effect. An input # set by the formula for this period still wins. - stored_input = self._input_set_meanwhile( - holder, period, input_state[2] - ) - return array if stored_input is None else stored_input + return self._uncached_fallback(holder, period, input_state) + array = holder.default_array() if variable.defined_for is not None: array = np.where(mask, array, variable.default_value) @@ -1999,6 +2005,18 @@ def _calculation_start(self) -> Tuple[int, int, int]: and inputs were set here (``_inputs_set``), and a sequence number.""" return self._input_epoch, self._inputs_set, next_sequence_number() + def _uncached_fallback( + self, holder: Holder, period: Period, input_state: Tuple[int, int, int] + ) -> ArrayLike: + """Return a default that must stay uncached even after convergence.""" + stored_input = self._input_set_meanwhile(holder, period, input_state[2]) + if stored_input is not None: + return stored_input + frames = _calculation_frames.get() + if frames and frames[-1].simulation is self: + frames[-1].uncached_fallbacks.add((holder.variable.name, period)) + return holder.default_array() + def _input_set_meanwhile( self, holder: Holder, period: Period, started_at: int ) -> Optional[ArrayLike]: diff --git a/tests/core/test_branch_input_fixed_point.py b/tests/core/test_branch_input_fixed_point.py index b51f4964e..9fc78ff79 100644 --- a/tests/core/test_branch_input_fixed_point.py +++ b/tests/core/test_branch_input_fixed_point.py @@ -10,6 +10,8 @@ import contextvars import threading +from types import SimpleNamespace +from unittest.mock import Mock import numpy as np import pytest @@ -164,6 +166,117 @@ def _input_first_result(values): return branch.calculate("fp_capped", "2020") +def _fallback_simulation(cache_mode="memory"): + """The review's branch reset, whose formula selects cross-period fallback.""" + calls = [] + + def cap(person, period): + return np.full(person.count, 1000.0) + + def work(person, period): + return np.full(person.count, 5.0) + + def carried(person, period): + calls.append(str(period)) + person("cap", period) + branch = person.simulation.get_branch("worker") + branch.delete_arrays("cap") + maximum = branch.calculate("work", period) + branch.set_input("cap", period, maximum) + # Returning None selects the default or a real input to carry. + + system = CountryTaxBenefitSystem() + system.auto_carry_over_input_variables = True + system.add_variables( + _variable("cap", periods.YEAR, cap), + _variable("work", periods.YEAR, work), + _variable("carried", periods.YEAR, carried), + ) + simulation = SimulationBuilder().build_default_simulation(system) + if cache_mode == "disk": + simulation.memory_config = MemoryConfig(0) + for population in simulation.populations.values(): + population._holders = {} + elif cache_mode == "drop": + simulation.memory_config = MemoryConfig(1, variables_to_drop=["carried"]) + for population in simulation.populations.values(): + population._holders = {} + elif cache_mode == "blacklist": + system.cache_blacklist = ["carried"] + simulation.opt_out_cache = True + return simulation, calls + + +@pytest.mark.parametrize("cache_mode", ["memory", "disk", "drop", "blacklist"]) +@pytest.mark.parametrize("future_input", [False, True]) +def test_fixed_point_fallback_does_not_hide_later_input(cache_mode, future_input): + simulation, calls = _fallback_simulation(cache_mode) + if future_input: + simulation.set_input("carried", "2022", [999]) + + assert simulation.calculate("carried", "2020").tolist() == [0.0] + assert len(calls) == 2 + + simulation.set_input("carried", "2019", [123]) + assert simulation.calculate("carried", "2020").tolist() == [123.0] + + +def test_fixed_point_fallback_is_absent_from_every_cache(monkeypatch, tmp_path): + simulation, calls = _fallback_simulation() + simulation.dataset = SimpleNamespace(file_path=tmp_path / "dataset.h5", name="test") + cache = Mock() + cache.get_cache_path.return_value = tmp_path / "missing-cache.h5" + monkeypatch.setattr(simulation_module, "SimulationMacroCache", lambda system: cache) + monkeypatch.setattr( + simulation, "check_macro_cache", lambda variable, period: variable == "carried" + ) + + assert simulation.calculate("carried", "2020").tolist() == [0.0] + assert len(calls) == 2 + assert simulation.get_array("carried", "2020") is None + assert ("carried", periods.period("2020")) not in simulation._fast_cache + cache.set_cache_value.assert_not_called() + + simulation.set_input("carried", "2019", [123]) + assert simulation.calculate("carried", "2020").tolist() == [123.0] + assert simulation.get_array("carried", "2020").tolist() == [123.0] + cache.set_cache_value.assert_called_once() + + +def test_formula_using_uncached_fallback_still_keeps_its_fixed_point(): + simulation, calls = _fallback_simulation() + + def outer(person, period): + return person("carried", period) + 1 + + simulation.tax_benefit_system.add_variable( + _variable("fp_fallback_outer", periods.YEAR, outer) + ) + assert simulation.calculate("fp_fallback_outer", "2020").tolist() == [1.0] + assert len(calls) == 2 + assert simulation.get_array("carried", "2020") is None + assert simulation.get_array("fp_fallback_outer", "2020").tolist() == [1.0] + + +def test_retry_can_cache_a_formula_result_after_an_uncached_fallback(): + simulation, calls = _fallback_simulation() + original = simulation.tax_benefit_system.get_variable("carried").get_formula("2020") + + def carried(person, period): + original(person, period) + if len(calls) > 1: + return np.full(person.count, 7.0) + + simulation.tax_benefit_system.update_variable( + _variable("carried", periods.YEAR, carried) + ) + assert simulation.calculate("carried", "2020").tolist() == [7.0] + assert len(calls) == 3 + assert simulation.get_array("carried", "2020").tolist() == [7.0] + assert simulation.calculate("carried", "2020").tolist() == [7.0] + assert len(calls) == 3 + + def _assert_settled(simulation, calls, expected, kept_in_holder=True): result = simulation.calculate("fp_outer", "2020") np.testing.assert_array_equal(result, expected) diff --git a/tests/core/test_branch_input_fixed_point_property.py b/tests/core/test_branch_input_fixed_point_property.py index 2d5358132..e1cedfce4 100644 --- a/tests/core/test_branch_input_fixed_point_property.py +++ b/tests/core/test_branch_input_fixed_point_property.py @@ -7,11 +7,54 @@ from tests.core.test_branch_input_fixed_point import ( _assert_settled, + _fallback_simulation, _input_first_result, _simulation, ) +@given( + operations=st.lists( + st.one_of( + st.tuples(st.just("set"), st.integers(2018, 2022), st.integers(-500, 500)), + st.tuples(st.just("calculate"), st.integers(2018, 2022), st.none()), + ), + min_size=1, + max_size=25, + ), + cache_mode=st.sampled_from(["memory", "disk", "drop", "blacklist"]), +) +@settings(max_examples=100, deadline=None) +def test_fixed_point_fallback_never_shadows_later_input(operations, cache_mode): + simulation, _ = _fallback_simulation(cache_mode) + inputs = {} + kept = {} + # Always include the reviewed calculate -> earlier input -> calculate + # witness, followed by arbitrary interleavings and repeated writes/reads. + operations = [ + ("calculate", 2020, None), + ("set", 2019, 123), + ("calculate", 2020, None), + *operations, + ] + for action, year, value in operations: + if action == "set": + simulation.set_input("carried", str(year), [value]) + inputs[year] = value + kept[year] = value + continue + + source = max((known for known in inputs if known <= year), default=None) + # Root inputs keep previous input-derived calculations. Defaults + # without a source must never acquire that persistent cache status. + if year not in kept and source is not None: + kept[year] = inputs[source] + expected = kept.get(year, 0) + assert simulation.calculate("carried", str(year)).tolist() == [expected] + if source is None: + assert simulation.get_array("carried", str(year)) is None + + @given( values=st.lists( st.floats(min_value=-500, max_value=500, allow_nan=False, width=32), From cbbd8c9d936cdd65d53a5a778ea0772cd725efd8 Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Thu, 22 Oct 2026 03:51:26 -0400 Subject: [PATCH 11/15] Check surviving derived storage metadata and portable periods --- ...t_disk_storage_delete_contained_periods.py | 25 ++--- .../core/test_disk_storage_delete_property.py | 93 +++++++++++++++---- 2 files changed, 81 insertions(+), 37 deletions(-) diff --git a/tests/core/test_disk_storage_delete_contained_periods.py b/tests/core/test_disk_storage_delete_contained_periods.py index 0505f0f99..8cdfc01ce 100644 --- a/tests/core/test_disk_storage_delete_contained_periods.py +++ b/tests/core/test_disk_storage_delete_contained_periods.py @@ -27,16 +27,6 @@ BACKENDS = ("memory", "disk") -# On Windows a ``:`` in a file name is invalid or names an alternate data -# stream, so disk storage cannot keep periods whose string form has one -# (``month:2025-01:3``, ``year:2025-03``). That predates this fix; tests of -# such periods run elsewhere. -skip_colon_periods_on_windows = pytest.mark.skipif( - os.name == "nt", - reason="On-disk keys for multi-unit periods contain ':', which Windows " - "file names cannot.", -) - @pytest.fixture def make_storage(tmp_path): @@ -195,7 +185,6 @@ def test_deleting_eternity_deletes_every_period_of_the_branch(make_storage, back assert known(storage) == [("reform", "2025-01")] -@skip_colon_periods_on_windows @pytest.mark.parametrize("backend", BACKENDS) def test_delete_multi_unit_periods(make_storage, backend): storage = make_storage(backend) @@ -237,9 +226,10 @@ def test_delete_multi_unit_periods(make_storage, backend): "2025-02-01", "2026-01", "ETERNITY", + "month:2025-01:2", + "year:2025-03", + "year:2024:2", ] -if os.name != "nt": - PERIODS += ["month:2025-01:2", "year:2025-03", "year:2024:2"] BRANCH_NAMES = ["default", "pre_tcja", "pre_tcja_ctc"] @@ -402,8 +392,6 @@ def test_put_rejects_unparseable_period_without_changing_values( @pytest.mark.parametrize("month", [1, 3]) def test_twelve_month_alias_remains_storable(make_storage, backend, month): period = periods.Period((periods.MONTH, periods.Instant((2025, month, 1)), 12)) - if os.name == "nt" and ":" in str(period): - pytest.skip("Windows file names cannot contain ':'") storage = make_storage(backend) storage.put(np.array([500.0]), period, "no_salt", derived=True) @@ -425,7 +413,8 @@ def test_twelve_month_alias_remains_storable(make_storage, backend, month): # is not a key ``put`` writes (no ``_``, no period after the last ``_``, or a # period not in its string form) names no value, and deleting any period of # any branch then raised parsing it, leaving the value it was to delete. -# Names without ``:``, which Windows file names cannot hold. +# The raw filenames below bypass storage's semicolon encoder. Only literal +# colon filenames remain excluded on Windows, which cannot create them. NOT_KEYS = [ "stray.npy", "default_garbage.npy", @@ -468,9 +457,9 @@ def test_restore_reads_back_every_key_without_warning(tmp_path): ("no_salt", "2025-03"), ("y_2019", "2025-03-05"), ("trailing_", "ETERNITY"), + ("pre_tcja", "year:2025-03"), + ("default", "month:2025-01:3"), ] - if os.name != "nt": - keys += [("pre_tcja", "year:2025-03"), ("default", "month:2025-01:3")] storage = OnDiskStorage(str(tmp_path), preserve_storage_dir=True) put_all(storage, keys) diff --git a/tests/core/test_disk_storage_delete_property.py b/tests/core/test_disk_storage_delete_property.py index 8f92bd761..06a884f58 100644 --- a/tests/core/test_disk_storage_delete_property.py +++ b/tests/core/test_disk_storage_delete_property.py @@ -10,8 +10,9 @@ removes the branch's values for every period that ``period`` contains, ``delete(None, branch_name)`` removes all of the branch's values, and eternal storage keeps one ``ETERNITY`` value per branch. After every step both -storages hold exactly the model's (branch, period) keys, with the values last -put there. +storages hold exactly the model's (branch, period) keys, with the values, +input/derived provenance and sequence numbers last put there. Deleting one +period must preserve all three for a derived value surviving in another. A second property restores a disk storage's directory, holding the files of random keys and of names that are not keys, and checks that ``restore`` reads @@ -58,14 +59,10 @@ def period_strategy(draw): return periods.period(periods.ETERNITY) unit = draw(st.sampled_from([periods.YEAR, periods.MONTH, periods.DAY])) year = draw(st.one_of(st.sampled_from([999, 1000]), st.integers(2023, 2026))) - if os.name == "nt": - # Windows file names cannot contain the ``:`` of a multi-unit or - # rolling-year period's string form. - month = 1 if unit == periods.YEAR else draw(st.integers(1, 12)) - size = 1 - else: - month = draw(st.integers(1, 12)) - size = draw(st.sampled_from([1, 2, 3, 12])) + # Disk filenames encode compound periods' colons as semicolons, so these + # periods are supported on Windows too. + month = draw(st.integers(1, 12)) + size = draw(st.sampled_from([1, 2, 3, 12])) # Both storages reject years and months anchored mid-month, whose string # form drops the day (before writing a file, so on Windows too). if unit == periods.DAY: @@ -86,13 +83,21 @@ def storable(branch_name, period): return ":" not in branch_name and parsed.start == periods.period(period).start -def put_both(memory, disk, value, period, branch_name): +def put_both( + memory, disk, value, period, branch_name, *, derived=False, sequence_number=None +): """Put ``value`` in both storages; whether both stored it. Both store it, or both raise the same ``ValueError``, as the model says.""" errors = [] for storage in (memory, disk): try: - storage.put(np.array([float(value)]), period, branch_name) + storage.put( + np.array([float(value)]), + period, + branch_name, + derived=derived, + sequence_number=sequence_number, + ) errors.append(None) except ValueError as error: errors.append(str(error)) @@ -108,12 +113,14 @@ def put_both(memory, disk, value, period, branch_name): period_strategy(), st.sampled_from(BRANCH_NAMES), st.integers(0, 1_000), + st.booleans(), ) delete_step = st.tuples( st.just("delete"), st.one_of(st.none(), period_strategy()), st.sampled_from(BRANCH_NAMES), st.none(), + st.none(), ) @@ -125,11 +132,31 @@ def stored(storage): } +def provenance(storage): + """Map each stored key to its derived mark and original sequence number.""" + separator = ":" if isinstance(storage, InMemoryStorage) else "_" + return { + (branch_name, str(period)): ( + storage.is_derived(period, branch_name), + storage._sequence_numbers[f"{branch_name}{separator}{period}"], + ) + for branch_name, period in storage.get_known_branch_periods() + } + + @hypothesis.settings( max_examples=200, deadline=None, suppress_health_check=[hypothesis.HealthCheck.too_slow], ) +@hypothesis.example( + is_eternal=False, + steps=[ + ("put", periods.period(2025), "default", 7, False), + ("put", periods.period(2026), "default", 9, True), + ("delete", periods.period(2025), "default", None, None), + ], +) @hypothesis.given( is_eternal=st.booleans(), steps=st.lists(st.one_of(put_step, put_step, delete_step), max_size=25), @@ -141,12 +168,26 @@ def test_disk_and_memory_storage_delete_the_same_keys(is_eternal, steps): disk = OnDiskStorage( directory, is_eternal=is_eternal, preserve_storage_dir=True ) - for action, period, branch_name, value in steps: + for sequence_number, (action, period, branch_name, value, derived) in enumerate( + steps, start=1 + ): if period is not None and is_eternal: period = periods.period(periods.ETERNITY) if action == "put": - if put_both(memory, disk, value, period, branch_name): - model[(branch_name, str(period))] = float(value) + if put_both( + memory, + disk, + value, + period, + branch_name, + derived=derived, + sequence_number=sequence_number, + ): + model[(branch_name, str(period))] = ( + float(value), + derived, + sequence_number, + ) else: memory.delete(period, branch_name) disk.delete(period, branch_name) @@ -160,7 +201,19 @@ def test_disk_and_memory_storage_delete_the_same_keys(is_eternal, steps): ) } - assert stored(disk) == stored(memory) == model + assert ( + stored(disk) + == stored(memory) + == {key: value for key, (value, _, _) in model.items()} + ) + assert ( + provenance(disk) + == provenance(memory) + == { + key: (derived, sequence_number) + for key, (_, derived, sequence_number) in model.items() + } + ) # Names of ``.npy`` files that are not keys: no ``_``, no period after the last @@ -168,6 +221,8 @@ def test_disk_and_memory_storage_delete_the_same_keys(is_eternal, steps): # case-insensitive file system that is the ``ETERNITY`` key's file.) NON_CANONICAL_PERIODS = ["2025-3", "2025-03-1", "999", "x", ""] if os.name != "nt": + # These are deliberately malformed raw filenames, written without the + # storage encoder. Windows cannot create literal-colon filenames. NON_CANONICAL_PERIODS += ["year:2025", "month:2025-01:12", "day:2025-03-01"] not_key_names = st.one_of( st.text(alphabet="abcxyz0123456789-", min_size=1, max_size=8), @@ -201,8 +256,8 @@ def test_restore_reads_back_the_keys_and_only_the_keys( disk = OnDiskStorage( directory, is_eternal=is_eternal, preserve_storage_dir=True ) - for _, period, branch_name, value in puts: - put_both(memory, disk, value, period, branch_name) + for _, period, branch_name, value, derived in puts: + put_both(memory, disk, value, period, branch_name, derived=derived) for name in not_keys: np.save(os.path.join(directory, name), np.array([-1.0])) @@ -221,7 +276,7 @@ def test_restore_reads_back_the_keys_and_only_the_keys( else: assert caught == [] - for _, period, branch_name, _ in deletes: + for _, period, branch_name, _, _ in deletes: memory.delete(period, branch_name) restored.delete(period, branch_name) assert stored(restored) == stored(memory) From e02a77b69855546fe93336fcb7349f36f5c25126 Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Thu, 22 Oct 2026 04:28:13 -0400 Subject: [PATCH 12/15] Exercise arbitrary fallback input and calculation sequences --- .../test_branch_input_fixed_point_property.py | 18 +++++++++--------- 1 file changed, 9 insertions(+), 9 deletions(-) diff --git a/tests/core/test_branch_input_fixed_point_property.py b/tests/core/test_branch_input_fixed_point_property.py index e1cedfce4..1bd27c814 100644 --- a/tests/core/test_branch_input_fixed_point_property.py +++ b/tests/core/test_branch_input_fixed_point_property.py @@ -3,7 +3,7 @@ import pytest pytest.importorskip("hypothesis") -from hypothesis import given, settings, strategies as st +from hypothesis import example, given, settings, strategies as st from tests.core.test_branch_input_fixed_point import ( _assert_settled, @@ -13,6 +13,14 @@ ) +@example( + operations=[ + ("calculate", 2020, None), + ("set", 2019, 123), + ("calculate", 2020, None), + ], + cache_mode="memory", +) @given( operations=st.lists( st.one_of( @@ -29,14 +37,6 @@ def test_fixed_point_fallback_never_shadows_later_input(operations, cache_mode): simulation, _ = _fallback_simulation(cache_mode) inputs = {} kept = {} - # Always include the reviewed calculate -> earlier input -> calculate - # witness, followed by arbitrary interleavings and repeated writes/reads. - operations = [ - ("calculate", 2020, None), - ("set", 2019, 123), - ("calculate", 2020, None), - *operations, - ] for action, year, value in operations: if action == "set": simulation.set_input("carried", str(year), [value]) From 551f935820078ae3c254a8e00924fc28c9b2f1fd Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Thu, 22 Oct 2026 23:08:08 -0400 Subject: [PATCH 13/15] Keep converted root fallbacks out of calculation caches --- .../fix-uprating-branch-visibility.fixed.md | 2 +- docs/usage/branches.md | 12 +- policyengine_core/simulations/simulation.py | 73 +++++++++-- tests/core/test_branch_input_fixed_point.py | 118 +++++++++++++++++- .../test_branch_input_fixed_point_property.py | 65 ++++++++-- 5 files changed, 236 insertions(+), 34 deletions(-) diff --git a/changelog.d/fix-uprating-branch-visibility.fixed.md b/changelog.d/fix-uprating-branch-visibility.fixed.md index 2abb0cac1..03b10fb3e 100644 --- a/changelog.d/fix-uprating-branch-visibility.fixed.md +++ b/changelog.d/fix-uprating-branch-visibility.fixed.md @@ -1 +1 @@ -Scope known periods and dumps to the branch's readable values while preserving derived marks, leave root uprating and carry-over fallback defaults uncached even after branch mutations converge so later inputs take effect, parse underscore-containing disk branch names safely, delete contained disk periods consistently with memory, reject ambiguous or unparseable storage keys before writing, and warn and skip malformed restored filenames. +Scope known periods and dumps to the branch's readable values while preserving derived marks, leave root uprating and carry-over fallback defaults uncached through period conversions and even after branch mutations converge so later inputs take effect, parse underscore-containing disk branch names safely, delete contained disk periods consistently with memory, reject ambiguous or unparseable storage keys before writing, and warn and skip malformed restored filenames. diff --git a/docs/usage/branches.md b/docs/usage/branches.md index d4dda2a5d..369ae23c6 100644 --- a/docs/usage/branches.md +++ b/docs/usage/branches.md @@ -239,8 +239,10 @@ depend on the input is calculated again) but not less, within these limits: as described above, and results calculated while an input is being set do not enter the fast cache. When uprating or automatic carry-over finds no source input, the root returns the fallback default without caching it, including - when repeated branch mutations converge to a fixed point. Setting an earlier - input can therefore change the next calculation of that period. + when repeated branch mutations converge to a fixed point. Period conversion + preserves this rule: a divided, summed or stock-selected value containing + such a fallback also stays uncached. Setting an earlier input can therefore + change the next calculation of the requested period. Branches can cache these defaults because later inputs invalidate them. - **Existing child branches keep their values.** An input set on a branch does not reach branches already created from it. @@ -280,9 +282,9 @@ say what each value was calculated from, nor what was read without being kept, so the restored simulation records that values from that number on may depend on anything: an input set for any variable on a branch of it drops all of them. A dump written before calculated values were recorded is restored with every -value as an input, as before, so such values never drop. With disk storage, a branch whose -name contains `_` cannot be dumped yet: `OnDiskStorage.get_known_periods` -splits its keys on every `_` (as on master). +value as an input, as before, so such values never drop. Disk storage splits +keys at the last underscore, so branch names containing `_` also round-trip +through dumps and restoration. Two related behaviours: a branch stops reading macro-cache files once an input is set on it, even one it already reads, since they are keyed by branch name and diff --git a/policyengine_core/simulations/simulation.py b/policyengine_core/simulations/simulation.py index c9e598944..ba48b6b74 100644 --- a/policyengine_core/simulations/simulation.py +++ b/policyengine_core/simulations/simulation.py @@ -1664,9 +1664,7 @@ def _calculate(self, variable_name: str, period: Period = None) -> ArrayLike: raise ValueError(f"Variable {variable_name} does not exist.") input_state = self._calculation_start() frames = _calculation_frames.get() - frame = frames[-1] if frames and frames[-1].simulation is self else None - if frame is not None: - frame.uncached_fallbacks.discard((variable_name, period)) + self._set_fallback_cache_status(variable_name, period, uncached=False) # Lets a custom ``set_input`` handler tell whether it calculated. self._calculations_started = self._calculations_started + 1 population = self.get_variable_population(variable_name) @@ -1750,7 +1748,13 @@ def _calculate(self, variable_name: str, period: Period = None) -> ArrayLike: if variable.definition_period == MONTH and period.unit == YEAR: if variable.quantity_type == QuantityType.STOCK: contained_months = period.get_subperiods(MONTH) - values = self._calculate(variable_name, contained_months[-1]) + source = contained_months[-1] + values = self._calculate(variable_name, source) + self._set_fallback_cache_status( + variable_name, + period, + uncached=self._fallback_is_uncached(variable_name, source), + ) else: values = self.calculate_add(variable_name, period) alternate_period_handling = True @@ -1758,11 +1762,22 @@ def _calculate(self, variable_name: str, period: Period = None) -> ArrayLike: alternate_period_handling = True if variable.quantity_type == QuantityType.STOCK: values = self._calculate(variable_name, period.this_year) + self._set_fallback_cache_status( + variable_name, + period, + uncached=self._fallback_is_uncached( + variable_name, period.this_year + ), + ) else: values = self.calculate_divide(variable_name, period) if alternate_period_handling: - if is_cache_available and self._may_keep(input_state): + if ( + is_cache_available + and not self._fallback_is_uncached(variable_name, period) + and self._may_keep(input_state) + ): smc.set_cache_value(cache_path, values) return values @@ -2012,11 +2027,30 @@ def _uncached_fallback( stored_input = self._input_set_meanwhile(holder, period, input_state[2]) if stored_input is not None: return stored_input - frames = _calculation_frames.get() - if frames and frames[-1].simulation is self: - frames[-1].uncached_fallbacks.add((holder.variable.name, period)) + self._set_fallback_cache_status(holder.variable.name, period, uncached=True) return holder.default_array() + def _set_fallback_cache_status( + self, variable_name: str, period: Period, *, uncached: bool + ) -> None: + # A conversion helper or a stale caller can have its own frame. Keep + # the returned value's eligibility visible after a nested frame exits, + # including when a retry replaces a fallback with a real result. + key = (variable_name, period) + for frame in _calculation_frames.get(): + if frame.simulation is self: + if uncached: + frame.uncached_fallbacks.add(key) + else: + frame.uncached_fallbacks.discard(key) + + def _fallback_is_uncached(self, variable_name: str, period: Period) -> bool: + key = (variable_name, period) + return any( + frame.simulation is self and key in frame.uncached_fallbacks + for frame in _calculation_frames.get() + ) + def _input_set_meanwhile( self, holder: Holder, period: Period, started_at: int ) -> Optional[ArrayLike]: @@ -2124,9 +2158,15 @@ def calculate_add( sub_periods = list(period.get_subperiods(variable.definition_period)) def total(): - return sum( - self.calculate(variable_name, sub_period) for sub_period in sub_periods - ) + result = 0 + uncached = False + for sub_period in sub_periods: + result = result + self.calculate(variable_name, sub_period) + # Preserve each returned term's status: a later term can + # recalculate an earlier period without changing this sum. + uncached |= self._fallback_is_uncached(variable_name, sub_period) + self._set_fallback_cache_status(variable_name, period, uncached=uncached) + return result # As in ``calculate``: only an outermost sum runs again, and while it # sums it is in flight, so its terms do not run again on their own @@ -2188,6 +2228,13 @@ def calculate_divide( input_state = self._calculation_start() computation_period = period.this_year result = self.calculate(variable_name, period=computation_period) / 12.0 + self._set_fallback_cache_status( + variable_name, + period, + uncached=self._fallback_is_uncached( + variable_name, computation_period + ), + ) return self._cache_option_result(variable, period, result, input_state) elif period.unit == periods.YEAR: return self.calculate(variable_name, period) @@ -2240,7 +2287,9 @@ def _cache_option_result( and raw holder inputs need not be registered there. Accepted cache writes never register the aggregate as a supplied input. """ - if variable.quantity_type == QuantityType.STOCK: + if variable.quantity_type == QuantityType.STOCK or self._fallback_is_uncached( + variable.name, period + ): return result routed = (variable.definition_period == MONTH and period.unit == YEAR) or ( variable.definition_period == YEAR and period.unit == MONTH diff --git a/tests/core/test_branch_input_fixed_point.py b/tests/core/test_branch_input_fixed_point.py index 9fc78ff79..80e9a7f13 100644 --- a/tests/core/test_branch_input_fixed_point.py +++ b/tests/core/test_branch_input_fixed_point.py @@ -21,7 +21,7 @@ from policyengine_core.experimental import MemoryConfig from policyengine_core.simulations import SimulationBuilder import policyengine_core.simulations.simulation as simulation_module -from policyengine_core.variables import Variable +from policyengine_core.variables import QuantityType, Variable def _variable(name, definition_period, formula=None): @@ -166,7 +166,7 @@ def _input_first_result(values): return branch.calculate("fp_capped", "2020") -def _fallback_simulation(cache_mode="memory"): +def _fallback_simulation(cache_mode="memory", definition_period=periods.YEAR): """The review's branch reset, whose formula selects cross-period fallback.""" calls = [] @@ -188,9 +188,9 @@ def carried(person, period): system = CountryTaxBenefitSystem() system.auto_carry_over_input_variables = True system.add_variables( - _variable("cap", periods.YEAR, cap), - _variable("work", periods.YEAR, work), - _variable("carried", periods.YEAR, carried), + _variable("cap", definition_period, cap), + _variable("work", definition_period, work), + _variable("carried", definition_period, carried), ) simulation = SimulationBuilder().build_default_simulation(system) if cache_mode == "disk": @@ -243,6 +243,73 @@ def test_fixed_point_fallback_is_absent_from_every_cache(monkeypatch, tmp_path): cache.set_cache_value.assert_called_once() +@pytest.mark.parametrize("cache_mode", ["memory", "disk", "drop", "blacklist"]) +@pytest.mark.parametrize( + "definition_period,quantity_type,method,requested,source,expected", + [ + (periods.YEAR, QuantityType.FLOW, "calculate", "2020-01", "2019", 10.25), + (periods.YEAR, QuantityType.FLOW, "calculate_divide", "2020-01", "2019", 10.25), + (periods.YEAR, QuantityType.STOCK, "calculate", "2020-01", "2019", 123), + (periods.MONTH, QuantityType.FLOW, "calculate", "2020", "2019-12", 1476), + ( + periods.MONTH, + QuantityType.FLOW, + "calculate_add", + "year:2020:2", + "2019-12", + 2952, + ), + (periods.MONTH, QuantityType.STOCK, "calculate", "2020", "2019-12", 123), + ], +) +def test_converted_fixed_point_fallback_stays_out_of_every_cache( + cache_mode, + definition_period, + quantity_type, + method, + requested, + source, + expected, + monkeypatch, + tmp_path, +): + simulation, _ = _fallback_simulation(cache_mode, definition_period) + simulation.tax_benefit_system.get_variable("carried").quantity_type = quantity_type + simulation.dataset = SimpleNamespace(file_path=tmp_path / "dataset.h5", name="test") + cache = Mock() + cache.get_cache_path.return_value = tmp_path / "missing-cache.h5" + monkeypatch.setattr(simulation_module, "SimulationMacroCache", lambda system: cache) + monkeypatch.setattr( + simulation, "check_macro_cache", lambda variable, period: variable == "carried" + ) + + assert getattr(simulation, method)("carried", requested).tolist() == [0.0] + holder = simulation.get_holder("carried") + assert holder.get_known_periods(simulation.branch_name) == [] + assert not any(name == "carried" for name, _ in simulation._fast_cache) + cache.set_cache_value.assert_not_called() + assert not simulation._open_frames + + simulation.set_input("carried", source, [123]) + fresh, _ = _fallback_simulation(cache_mode, definition_period) + fresh.tax_benefit_system.get_variable("carried").quantity_type = quantity_type + fresh.set_input("carried", source, [123]) + assert getattr(fresh, method)("carried", requested).tolist() == [expected] + assert getattr(simulation, method)("carried", requested).tolist() == [expected] + + +@pytest.mark.parametrize("method", ["calculate", "calculate_add"]) +def test_sum_with_some_real_inputs_does_not_cache_its_fallback_terms(method): + simulation, _ = _fallback_simulation(definition_period=periods.MONTH) + simulation.set_input("carried", "2020-06", [12]) + assert getattr(simulation, method)("carried", "2020").tolist() == [84.0] + assert simulation.get_array("carried", "2020") is None + assert ("carried", periods.period("2020")) not in simulation._fast_cache + + simulation.set_input("carried", "2019-12", [6]) + assert getattr(simulation, method)("carried", "2020").tolist() == [114.0] + + def test_formula_using_uncached_fallback_still_keeps_its_fixed_point(): simulation, calls = _fallback_simulation() @@ -277,6 +344,47 @@ def carried(person, period): assert len(calls) == 3 +@pytest.mark.parametrize("method", ["calculate", "calculate_divide"]) +def test_converted_retry_can_cache_a_real_result_after_fallback(method): + simulation, calls = _fallback_simulation() + original = simulation.tax_benefit_system.get_variable("carried").get_formula("2020") + + def carried(person, period): + original(person, period) + if len(calls) > 1: + return np.full(person.count, 84.0) + + simulation.tax_benefit_system.update_variable( + _variable("carried", periods.YEAR, carried) + ) + assert getattr(simulation, method)("carried", "2020-01").tolist() == [7.0] + assert len(calls) == 3 + assert simulation.get_array("carried", "2020-01").tolist() == [7.0] + assert simulation.calculate("carried", "2020-01").tolist() == [7.0] + assert len(calls) == 3 + + +@pytest.mark.parametrize("method", ["calculate", "calculate_add"]) +def test_sum_retains_fallback_status_if_a_later_term_recalculates_its_source(method): + simulation, _ = _fallback_simulation(definition_period=periods.MONTH) + + def carried(person, period): + if period.start.month == 2: + # January's already-returned fallback remains a term of the sum, + # even though its current calculation now reads a real input. + person.simulation.set_input("carried", "2020-01", [12]) + person.simulation.calculate("carried", "2020-01") + + simulation.tax_benefit_system.update_variable( + _variable("carried", periods.MONTH, carried) + ) + assert getattr(simulation, method)("carried", "2020").tolist() == [132.0] + assert simulation.get_array("carried", "2020") is None + assert ("carried", periods.period("2020")) not in simulation._fast_cache + assert getattr(simulation, method)("carried", "2020").tolist() == [144.0] + assert simulation.get_array("carried", "2020").tolist() == [144.0] + + def _assert_settled(simulation, calls, expected, kept_in_holder=True): result = simulation.calculate("fp_outer", "2020") np.testing.assert_array_equal(result, expected) diff --git a/tests/core/test_branch_input_fixed_point_property.py b/tests/core/test_branch_input_fixed_point_property.py index 1bd27c814..fe084585f 100644 --- a/tests/core/test_branch_input_fixed_point_property.py +++ b/tests/core/test_branch_input_fixed_point_property.py @@ -5,6 +5,8 @@ pytest.importorskip("hypothesis") from hypothesis import example, given, settings, strategies as st +from policyengine_core import periods + from tests.core.test_branch_input_fixed_point import ( _assert_settled, _fallback_simulation, @@ -21,11 +23,37 @@ ], cache_mode="memory", ) +@example( + operations=[ + ("calculate", "2020-01", None), + ("set", 2019, 123), + ("calculate", "2020-01", None), + ], + cache_mode="memory", +) +@example( + operations=[ + ("set", 2018, 0), + ("calculate", "2018-01", None), + ("set", 2018, 1), + ("calculate", "2018-01", None), + ], + cache_mode="memory", +) @given( operations=st.lists( st.one_of( st.tuples(st.just("set"), st.integers(2018, 2022), st.integers(-500, 500)), - st.tuples(st.just("calculate"), st.integers(2018, 2022), st.none()), + st.tuples( + st.just("calculate"), + st.one_of( + st.integers(2018, 2022), + st.tuples(st.integers(2018, 2022), st.integers(1, 12)).map( + lambda value: f"{value[0]}-{value[1]:02d}" + ), + ), + st.none(), + ), ), min_size=1, max_size=25, @@ -37,22 +65,37 @@ def test_fixed_point_fallback_never_shadows_later_input(operations, cache_mode): simulation, _ = _fallback_simulation(cache_mode) inputs = {} kept = {} - for action, year, value in operations: + for action, requested, value in operations: + period = periods.period(requested) + year = period.start.year if action == "set": - simulation.set_input("carried", str(year), [value]) + simulation.set_input("carried", period, [value]) inputs[year] = value - kept[year] = value + # Writing an annual input removes overlapping calculated months + # of this variable, even on the root (independent of branch + # dependency invalidation). + kept = { + known: result + for known, result in kept.items() + if known.start.year != year + } + kept[period] = value continue source = max((known for known in inputs if known <= year), default=None) - # Root inputs keep previous input-derived calculations. Defaults - # without a source must never acquire that persistent cache status. - if year not in kept and source is not None: - kept[year] = inputs[source] - expected = kept.get(year, 0) - assert simulation.calculate("carried", str(year)).tolist() == [expected] + # Root inputs keep previous input-derived calculations, including a + # converted month. A monthly request does not itself retain the year + # while the branch reset makes its underlying calculation stale. + if period not in kept and source is not None: + annual = kept.get(period.this_year, inputs[source]) + kept[period] = annual / 12 if period.unit == periods.MONTH else annual + expected = kept.get(period, 0) + assert simulation.calculate("carried", period).tolist() == pytest.approx( + [expected] + ) if source is None: - assert simulation.get_array("carried", str(year)) is None + assert simulation.get_array("carried", period) is None + assert simulation.get_array("carried", period.this_year) is None @given( From 994e3ba0d7ddbc9baaca796e585e3359a133a197 Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Thu, 22 Oct 2026 23:08:16 -0400 Subject: [PATCH 14/15] Align cache fixtures with transient root fallback semantics --- tests/core/test_holder_write_fast_cache.py | 24 +++++++++++++++------ tests/core/test_on_disk_storage_clones.py | 4 +++- tests/core/test_restore_input_registry.py | 3 ++- tests/core/test_subsample_deleted_inputs.py | 16 +++++++++----- 4 files changed, 34 insertions(+), 13 deletions(-) diff --git a/tests/core/test_holder_write_fast_cache.py b/tests/core/test_holder_write_fast_cache.py index f928f04bb..3c53f2fc1 100644 --- a/tests/core/test_holder_write_fast_cache.py +++ b/tests/core/test_holder_write_fast_cache.py @@ -20,6 +20,7 @@ from policyengine_core import periods from policyengine_core.experimental import MemoryConfig +from policyengine_core.variables import Variable from tests.fixtures.uprated_inputs import build_simulation, build_system @@ -208,14 +209,25 @@ def test_a_write_to_disk_storage_replaces_a_calculated_value(system): def test_a_write_replaces_an_eternity_value_requested_without_a_period(system): - # ``calculate`` with no period caches the result under ``None``. - simulation = build_simulation(system) - assert simulation.calculate("eternal_code").tolist() == [0, 0] - assert ("eternal_code", None) in simulation._fast_cache + class eternal_code_plus_one(Variable): + def formula(person, period): + # The top-level request intentionally has no period. Population + # lookups inside formulas still require an explicit one. + return person("eternal_code", periods.ETERNITY) + 1 + + local_system = system.clone() + local_system.update_variable(eternal_code_plus_one) + # Calculate a formula over a real input: supplied inputs bypass the fast + # cache, while a missing root's default must stay uncached. + simulation = build_simulation(local_system, [("eternal_code", "2012", [0, 0])]) + assert simulation.calculate("eternal_code_plus_one").tolist() == [1, 1] + assert ("eternal_code_plus_one", None) in simulation._fast_cache - simulation.get_holder("eternal_code").set_input(periods.period("2015"), [8, 9]) + simulation.get_holder("eternal_code_plus_one").set_input( + periods.period("2015"), [8, 9] + ) - assert simulation.calculate("eternal_code").tolist() == [8, 9] + assert simulation.calculate("eternal_code_plus_one").tolist() == [8, 9] def _disk_simulation(system): diff --git a/tests/core/test_on_disk_storage_clones.py b/tests/core/test_on_disk_storage_clones.py index 911a7d3e2..cf348530c 100644 --- a/tests/core/test_on_disk_storage_clones.py +++ b/tests/core/test_on_disk_storage_clones.py @@ -62,7 +62,9 @@ def _stored(holder, period): @pytest.mark.parametrize("variable", list(HELPER_INPUTS), ids=["dispatch", "divide"]) def test_helper_input_on_a_clone_leaves_the_source_unchanged(system, variable): - parent = simulation(system, {}) + # A real uprating source makes December a stored derived value under + # both carry-over settings; a missing root input stays uncached. + parent = simulation(system, {variable: {"2011-12": [1.0, 2.0]}}) holder = _on_disk(parent, variable) calculated = parent.calculate(variable, "2012-12").copy() assert holder._disk_storage.is_derived(periods.period("2012-12")) diff --git a/tests/core/test_restore_input_registry.py b/tests/core/test_restore_input_registry.py index dbaa825ff..ecf429c71 100644 --- a/tests/core/test_restore_input_registry.py +++ b/tests/core/test_restore_input_registry.py @@ -96,6 +96,7 @@ def test_dump_lists_the_calculated_periods_and_marks_its_format(system, tmp_path [ ("uprated_count", "2012", [1001, 77]), ("uprated_count", "2014", [2001, 177]), + ("uprated_amount", "2012", [100.0, 50.0]), ], ) simulation.calculate("uprated_count", "2013") @@ -107,7 +108,7 @@ def calculated(variable): return sorted(path.read_text().split()) if path.exists() else [] assert calculated("uprated_count") == ["2013"] - # Calculated from the defaulted input, which is calculated too. + # Both the uprated value and the formula reading it are derived. assert calculated("doubled_amount") == ["2013"] assert calculated("uprated_amount") == ["2013"] assert (tmp_path / DUMP_FORMAT_FILE).exists() diff --git a/tests/core/test_subsample_deleted_inputs.py b/tests/core/test_subsample_deleted_inputs.py index e6b276673..5dbbe5d31 100644 --- a/tests/core/test_subsample_deleted_inputs.py +++ b/tests/core/test_subsample_deleted_inputs.py @@ -70,16 +70,22 @@ def delete_override(simulation, variable, through_holder): @pytest.mark.parametrize("variable", ["doubled", "deleted_leaf", "deleted_eternal"]) @pytest.mark.parametrize("on_disk", [False, True], ids=["memory", "disk"]) @pytest.mark.parametrize("through_holder", [False, True], ids=["simulation", "holder"]) -def test_deleted_input_recalculated_as_derived_is_not_exportable( +def test_deleted_input_recalculation_is_not_exportable( variable, on_disk, through_holder ): simulation = prepare(variable, on_disk) delete_override(simulation, variable, through_holder) - simulation.calculate(variable, DATASET_YEAR) + result = simulation.calculate(variable, DATASET_YEAR) - assert simulation.get_holder(variable).is_derived( - period(DATASET_YEAR), simulation.branch_name - ) + holder = simulation.get_holder(variable) + if variable == "deleted_leaf": + # A missing input-only root returns its default without storing it. + np.testing.assert_array_equal(result, np.full(simulation.persons.count, 17)) + assert holder.get_array(period(DATASET_YEAR), simulation.branch_name) is None + assert (variable, period(DATASET_YEAR)) not in simulation._fast_cache + else: + # Formula results remain stored as derived values, never new inputs. + assert holder.is_derived(period(DATASET_YEAR), simulation.branch_name) assert simulation._get_set_input_periods(variable) == [] From 6ed1b402622db9ae49bfa40b0993e89e5e805969 Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Fri, 23 Oct 2026 01:47:33 -0400 Subject: [PATCH 15/15] Preserve inputs supplied during period conversions Resolve requested-period inputs before conversion cache guards, and retain the full ADD input window across retries. Cover both review witnesses, STOCK and integer conversion boundaries, and one-time formula inputs with Hypothesis. Correct the existing sequence model for legitimate retention of source-year results. --- .../fix-uprating-branch-visibility.fixed.md | 2 +- docs/usage/branches.md | 6 +- policyengine_core/simulations/simulation.py | 31 ++- tests/core/test_branch_input_fixed_point.py | 176 ++++++++++++++++++ .../test_branch_input_fixed_point_property.py | 127 ++++++++++++- 5 files changed, 329 insertions(+), 13 deletions(-) diff --git a/changelog.d/fix-uprating-branch-visibility.fixed.md b/changelog.d/fix-uprating-branch-visibility.fixed.md index 03b10fb3e..a2c5c9e58 100644 --- a/changelog.d/fix-uprating-branch-visibility.fixed.md +++ b/changelog.d/fix-uprating-branch-visibility.fixed.md @@ -1 +1 @@ -Scope known periods and dumps to the branch's readable values while preserving derived marks, leave root uprating and carry-over fallback defaults uncached through period conversions and even after branch mutations converge so later inputs take effect, parse underscore-containing disk branch names safely, delete contained disk periods consistently with memory, reject ambiguous or unparseable storage keys before writing, and warn and skip malformed restored filenames. +Scope known periods and dumps to the branch's readable values while preserving derived marks, leave root uprating and carry-over fallback defaults uncached through period conversions and even after branch mutations converge so later inputs take effect, return inputs supplied during conversion on the first call, parse underscore-containing disk branch names safely, delete contained disk periods consistently with memory, reject ambiguous or unparseable storage keys before writing, and warn and skip malformed restored filenames. diff --git a/docs/usage/branches.md b/docs/usage/branches.md index 369ae23c6..eec1ae1df 100644 --- a/docs/usage/branches.md +++ b/docs/usage/branches.md @@ -241,8 +241,10 @@ depend on the input is calculated again) but not less, within these limits: input, the root returns the fallback default without caching it, including when repeated branch mutations converge to a fixed point. Period conversion preserves this rule: a divided, summed or stock-selected value containing - such a fallback also stays uncached. Setting an earlier input can therefore - change the next calculation of the requested period. + such a fallback also stays uncached. If a formula supplies an input for the + requested period during conversion, that input is returned on the first + call, before deciding whether to cache the converted result. Setting an + earlier input can therefore change the next calculation of the requested period. Branches can cache these defaults because later inputs invalidate them. - **Existing child branches keep their values.** An input set on a branch does not reach branches already created from it. diff --git a/policyengine_core/simulations/simulation.py b/policyengine_core/simulations/simulation.py index ba48b6b74..d9bb9e8a7 100644 --- a/policyengine_core/simulations/simulation.py +++ b/policyengine_core/simulations/simulation.py @@ -1773,6 +1773,13 @@ def _calculate(self, variable_name: str, period: Period = None) -> ArrayLike: values = self.calculate_divide(variable_name, period) if alternate_period_handling: + # A formula for the source period can supply the requested + # period itself, including during STOCK selection. That input + # wins even when the converted fallback must stay uncached. + stored_input = self._input_set_meanwhile(holder, period, input_state[2]) + if stored_input is not None: + values = stored_input + self._set_fallback_cache_status(variable_name, period, uncached=False) if ( is_cache_available and not self._fallback_is_uncached(variable_name, period) @@ -2177,6 +2184,7 @@ def total(): try: with frame if frame is not None else nullcontext(): input_state = self._calculation_start() + started_at = input_state[2] result = total() previous_transition = None for _ in range(_RERUNS_AFTER_INPUT_CHANGE if outermost else 0): @@ -2192,7 +2200,9 @@ def total(): frame.restart() input_state = self._calculation_start() result = total() - return self._cache_option_result(variable, period, result, input_state) + return self._cache_option_result( + variable, period, result, input_state, started_at=started_at + ) finally: self._calculations_in_flight -= 1 @@ -2251,6 +2261,8 @@ def _cache_option_result( period: Period, result: ArrayLike, input_state: Tuple[int, int, int], + *, + started_at: Optional[int] = None, ) -> ArrayLike: """Cache an ADD or DIVIDE result at ``period`` if a plain read would return it. @@ -2286,7 +2298,20 @@ def _cache_option_result( and export; it can name a disk input hidden by a derived memory value, and raw holder inputs need not be registered there. Accepted cache writes never register the aggregate as a supplied input. + + ``started_at`` preserves the full sum's input window across retries; + ``input_state`` controls cache eligibility for the final attempt. """ + holder = self.get_holder(variable.name) + # Resolve formula-time inputs before any cache rejection: returning + # an uncached fallback (or a result of another dtype) must not hide + # an input supplied for the requested period during conversion. + stored_input = self._input_set_meanwhile( + holder, period, input_state[2] if started_at is None else started_at + ) + if stored_input is not None: + self._set_fallback_cache_status(variable.name, period, uncached=False) + return stored_input if variable.quantity_type == QuantityType.STOCK or self._fallback_is_uncached( variable.name, period ): @@ -2296,9 +2321,7 @@ def _cache_option_result( ) if not routed or np.asarray(result).dtype != variable.dtype: return result - return self._cache_result( - self.get_holder(variable.name), result, period, input_state - ) + return self._cache_result(holder, result, period, input_state) def calculate_output(self, variable_name: str, period: Period = None) -> ArrayLike: """ diff --git a/tests/core/test_branch_input_fixed_point.py b/tests/core/test_branch_input_fixed_point.py index 80e9a7f13..eb851991e 100644 --- a/tests/core/test_branch_input_fixed_point.py +++ b/tests/core/test_branch_input_fixed_point.py @@ -310,6 +310,182 @@ def test_sum_with_some_real_inputs_does_not_cache_its_fallback_terms(method): assert getattr(simulation, method)("carried", "2020").tolist() == [114.0] +@pytest.mark.parametrize("method", ["calculate", "calculate_divide"]) +def test_divided_fallback_yields_to_input_supplied_during_calculation(method): + class supplied_during_conversion(Variable): + value_type = float + entity = entities.Person + label = "Input supplied while dividing a fallback" + definition_period = periods.YEAR + quantity_type = QuantityType.FLOW + set_input = None + + def formula(person, period): + person.simulation.set_input("supplied_during_conversion", "2020-02", [123]) + return None + + system = CountryTaxBenefitSystem() + system.auto_carry_over_input_variables = True + system.add_variable(supplied_during_conversion) + simulation = SimulationBuilder().build_default_simulation(system) + + assert getattr(simulation, method)( + "supplied_during_conversion", "2020-02" + ).tolist() == [123.0] + holder = simulation.get_holder("supplied_during_conversion") + assert holder._is_input(periods.period("2020-02"), simulation.branch_name) + assert simulation.get_array("supplied_during_conversion", "2020-02").tolist() == [ + 123.0 + ] + assert simulation.get_array("supplied_during_conversion", "2020") is None + assert ("supplied_during_conversion", periods.period("2020")) not in ( + simulation._fast_cache + ) + assert not simulation._open_frames + + +@pytest.mark.parametrize("method", ["calculate", "calculate_add"]) +def test_summed_fallback_yields_to_input_supplied_during_calculation(method): + class supplied_during_conversion(Variable): + value_type = float + entity = entities.Person + label = "Input supplied while summing fallbacks" + definition_period = periods.MONTH + quantity_type = QuantityType.FLOW + set_input = None + + def formula(person, period): + if period.start.month == 6: + person.simulation.set_input("supplied_during_conversion", "2020", [120]) + return None + + system = CountryTaxBenefitSystem() + system.auto_carry_over_input_variables = True + system.add_variable(supplied_during_conversion) + simulation = SimulationBuilder().build_default_simulation(system) + + assert getattr(simulation, method)( + "supplied_during_conversion", "2020" + ).tolist() == [120.0] + holder = simulation.get_holder("supplied_during_conversion") + assert holder._is_input(periods.period("2020"), simulation.branch_name) + assert simulation.get_array("supplied_during_conversion", "2020").tolist() == [ + 120.0 + ] + for month in range(1, 6): + fallback_period = periods.period(f"2020-{month:02d}") + assert ( + simulation.get_array("supplied_during_conversion", fallback_period) is None + ) + assert ("supplied_during_conversion", fallback_period) not in ( + simulation._fast_cache + ) + assert not simulation._open_frames + + +@pytest.mark.parametrize( + "definition_period,quantity_type,value_type,method,formula_result", + [ + pytest.param( + periods.YEAR, QuantityType.STOCK, float, "calculate", None, id="stock-month" + ), + pytest.param( + periods.YEAR, + QuantityType.STOCK, + float, + "calculate_divide", + None, + id="stock-divide", + ), + pytest.param( + periods.MONTH, QuantityType.STOCK, float, "calculate", 7, id="stock-year" + ), + pytest.param( + periods.YEAR, + QuantityType.FLOW, + int, + "calculate", + None, + id="int-fallback-month", + ), + pytest.param( + periods.YEAR, + QuantityType.FLOW, + int, + "calculate_divide", + None, + id="int-fallback-divide", + ), + pytest.param( + periods.YEAR, + QuantityType.FLOW, + int, + "calculate", + 84, + id="int-formula-month", + ), + pytest.param( + periods.YEAR, + QuantityType.FLOW, + int, + "calculate_divide", + 84, + id="int-formula-divide", + ), + ], +) +def test_conversion_boundaries_preserve_input_supplied_during_calculation( + definition_period, quantity_type, value_type, method, formula_result +): + requested = "2020" if definition_period == periods.MONTH else "2020-02" + supplied = 120 if definition_period == periods.MONTH else 123 + + def formula(person, period): + if definition_period == periods.YEAR or period.start.month == 12: + person.simulation.set_input( + "supplied_at_conversion_boundary", requested, [supplied] + ) + # A real result distinguishes STOCK's selected December value and + # integer DIVIDE's floating result from the newly supplied input. + if formula_result is not None: + return np.full(person.count, formula_result) + return None + + system = CountryTaxBenefitSystem() + system.auto_carry_over_input_variables = True + system.add_variable( + type( + "supplied_at_conversion_boundary", + (Variable,), + { + "value_type": value_type, + "entity": entities.Person, + "label": "Input supplied at a period conversion boundary", + "definition_period": definition_period, + "quantity_type": quantity_type, + "set_input": None, + "formula": formula, + }, + ) + ) + simulation = SimulationBuilder().build_default_simulation(system) + + assert getattr(simulation, method)( + "supplied_at_conversion_boundary", requested + ).tolist() == [supplied] + holder = simulation.get_holder("supplied_at_conversion_boundary") + assert holder._is_input(periods.period(requested), simulation.branch_name) + assert simulation.get_array( + "supplied_at_conversion_boundary", requested + ).tolist() == [supplied] + if formula_result is None: + assert simulation.get_array("supplied_at_conversion_boundary", "2020") is None + assert ("supplied_at_conversion_boundary", periods.period("2020")) not in ( + simulation._fast_cache + ) + assert not simulation._open_frames + + def test_formula_using_uncached_fallback_still_keeps_its_fixed_point(): simulation, calls = _fallback_simulation() diff --git a/tests/core/test_branch_input_fixed_point_property.py b/tests/core/test_branch_input_fixed_point_property.py index fe084585f..90e236ebc 100644 --- a/tests/core/test_branch_input_fixed_point_property.py +++ b/tests/core/test_branch_input_fixed_point_property.py @@ -40,6 +40,16 @@ ], cache_mode="memory", ) +@example( + operations=[ + ("set", 2019, 1), + ("calculate", 2018, None), + ("calculate", "2020-01", None), + ("set", 2019, 0), + ("calculate", 2020, None), + ], + cache_mode="memory", +) @given( operations=st.lists( st.one_of( @@ -84,18 +94,123 @@ def test_fixed_point_fallback_never_shadows_later_input(operations, cache_mode): source = max((known for known in inputs if known <= year), default=None) # Root inputs keep previous input-derived calculations, including a - # converted month. A monthly request does not itself retain the year - # while the branch reset makes its underlying calculation stale. - if period not in kept and source is not None: - annual = kept.get(period.this_year, inputs[source]) - kept[period] = annual / 12 if period.unit == periods.MONTH else annual - expected = kept.get(period, 0) + # converted month or its source year. Derive each expected value from + # explicit inputs and previously modeled values, never cache contents. + annual = kept.get(period.this_year, inputs[source] if source is not None else 0) + converted = annual / 12 if period.unit == periods.MONTH else annual + expected = kept.get(period, converted) assert simulation.calculate("carried", period).tolist() == pytest.approx( [expected] ) + + # A reused worker can leave the source year fresh enough to cache; + # drop/blacklist modes can omit the converted month. Observe only + # whether caches retain the independently modeled values, validating + # every retained value before remembering its expected counterpart. + candidates = {period.this_year: annual, period: expected} + for candidate, modeled in candidates.items(): + stored = simulation.get_array("carried", candidate) + fast = simulation._fast_cache.get(("carried", candidate)) + for cached in (stored, fast): + if cached is not None: + assert cached.tolist() == pytest.approx([modeled]) + if stored is not None or fast is not None: + kept[candidate] = modeled + else: + kept.pop(candidate, None) if source is None: assert simulation.get_array("carried", period) is None assert simulation.get_array("carried", period.this_year) is None + assert ("carried", period) not in simulation._fast_cache + assert ("carried", period.this_year) not in simulation._fast_cache + + +@example( + definition_period=periods.YEAR, + use_option=False, + month=2, + value=123, + supplied_during_calculation=True, + cache_mode="memory", +) +@example( + definition_period=periods.MONTH, + use_option=False, + month=6, + value=120, + supplied_during_calculation=True, + cache_mode="memory", +) +@example( + definition_period=periods.MONTH, + use_option=True, + month=6, + value=120, + supplied_during_calculation=True, + cache_mode="memory", +) +@given( + definition_period=st.sampled_from([periods.YEAR, periods.MONTH]), + use_option=st.booleans(), + month=st.integers(2, 12), + value=st.integers(-500, 500), + supplied_during_calculation=st.booleans(), + cache_mode=st.sampled_from(["memory", "disk", "drop", "blacklist"]), +) +@settings(max_examples=100, deadline=None) +def test_converted_fallback_matches_input_supplied_during_calculation( + definition_period, + use_option, + month, + value, + supplied_during_calculation, + cache_mode, +): + simulation, _ = _fallback_simulation(cache_mode, definition_period) + variable = simulation.tax_benefit_system.get_variable("carried") + original = variable.get_formula("2020") + requested = f"2020-{month:02d}" if definition_period == periods.YEAR else "2020" + supplied = False + + def carried(person, period): + nonlocal supplied + original(person, period) + if ( + supplied_during_calculation + and not supplied + and (definition_period == periods.YEAR or period.start.month == month) + ): + person.simulation.set_input("carried", requested, [value]) + # Retrying the sum must still recognize this first-attempt input + # without relying on the formula to write it again. + supplied = True + # Returning None still selects fallback; a newly supplied input for + # the requested period takes precedence over the converted result. + + # Preserve the variable shared with any holders already constructed for + # disk/drop storage, including the input handler those holders consult. + variable.formulas["2020-01-01"] = carried + variable.set_input = None + option = ( + "calculate_divide" if definition_period == periods.YEAR else "calculate_add" + ) + method = option if use_option else "calculate" + + # The model depends only on whether the formula supplies the requested + # input. In particular, it never divides that input or sums carried terms. + expected = value if supplied_during_calculation else 0 + assert getattr(simulation, method)("carried", requested).tolist() == pytest.approx( + [expected] + ) + if supplied_during_calculation: + assert simulation.get_array("carried", requested).tolist() == [value] + else: + assert ( + simulation.get_holder("carried").get_known_periods(simulation.branch_name) + == [] + ) + assert not any(name == "carried" for name, _ in simulation._fast_cache) + assert not simulation._open_frames @given(