From bbb4f644019eef05611ff17dc9f1f4ef3ba545af Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Sun, 27 Sep 2026 15:50:27 -0400 Subject: [PATCH 1/7] 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 2/7] 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 3/7] 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 4/7] 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 5/7] 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 6/7] 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 7/7] 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(