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..268404cc9 --- /dev/null +++ b/changelog.d/fix-uprating-branch-visibility.fixed.md @@ -0,0 +1 @@ +Scope known periods and dumps to the branch's readable values while preserving derived marks, parse underscore-containing disk branch names safely, delete contained disk periods consistently with memory, reject ambiguous or unparseable storage keys before writing, and warn and skip malformed restored filenames. diff --git a/docs/python_api/data_storage.md b/docs/python_api/data_storage.md index 4307934f6..20b87a1e7 100644 --- a/docs/python_api/data_storage.md +++ b/docs/python_api/data_storage.md @@ -20,6 +20,20 @@ The `policyengine_core.data_storage` module contains the classes that are used t :show-inheritance: ``` +## 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: +``` + ## 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 3c5dbbb3f..bfbfafd5a 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 @@ -48,6 +50,41 @@ def _epoch() -> tuple: return os.getpid(), _forks +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 + + +def _is_key(key: str) -> bool: + """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 + 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.""" + 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 @@ -275,6 +312,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) @@ -328,12 +370,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 } self._derived.intersection_update(self._files) return @@ -342,27 +385,34 @@ 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) } self._derived.intersection_update(self._files) 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: """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 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 a new file instead, since the clone, copy or forked process reads @@ -374,14 +424,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..d17401b88 --- /dev/null +++ b/policyengine_core/data_storage/storage_keys.py @@ -0,0 +1,43 @@ +"""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`` 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 +from policyengine_core.periods import Period + + +def check_storable(branch_name: str, period: Period) -> None: + """Raise ``ValueError`` for an ambiguous branch or period serialization. + + 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. 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( + 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)." + ) + 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/policyengine_core/holders/holder.py b/policyengine_core/holders/holder.py index 324e3d7e9..54b7ae129 100644 --- a/policyengine_core/holders/holder.py +++ b/policyengine_core/holders/holder.py @@ -205,11 +205,25 @@ 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 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 = 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 123a5f1ae..0c708511d 100644 --- a/policyengine_core/simulations/simulation.py +++ b/policyengine_core/simulations/simulation.py @@ -1082,8 +1082,12 @@ 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 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 for known_period in known_periods diff --git a/policyengine_core/tools/simulation_dumper.py b/policyengine_core/tools/simulation_dumper.py index 65b3c6644..da665e863 100644 --- a/policyengine_core/tools/simulation_dumper.py +++ b/policyengine_core/tools/simulation_dumper.py @@ -70,13 +70,17 @@ def restore_simulation(directory, tax_benefit_system, **kwargs): 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 holder.get_known_periods(): - value = holder.get_array(period) - 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): + if holder.is_derived(period, branch_name): derived_periods.add(str(period)) if derived_periods: path = os.path.join(disk_storage.storage_dir, DERIVED_PERIODS_FILE) 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..0505f0f99 --- /dev/null +++ b/tests/core/test_disk_storage_delete_contained_periods.py @@ -0,0 +1,486 @@ +"""``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")] + + +# ----- 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", + ] + + +@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 +# 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", + "default_999.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 new file mode 100644 index 000000000..8f92bd761 --- /dev/null +++ b/tests/core/test_disk_storage_delete_property.py @@ -0,0 +1,227 @@ +"""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: +``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 ``_``. +``test_disk_storage_delete_contained_periods.py`` pins the same behaviour +with examples. +""" + +import os +import tempfile +import warnings + +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_", + # No key can hold a branch name containing ``:``. + "my:reform", +) + + +@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.one_of(st.sampled_from([999, 1000]), st.integers(2023, 2026))) + if os.name == "nt": + # Windows file names cannot contain the ``:`` of a multi-unit or + # rolling-year period's string form. + month = 1 if unit == periods.YEAR else draw(st.integers(1, 12)) + size = 1 + else: + month = draw(st.integers(1, 12)) + size = draw(st.sampled_from([1, 2, 3, 12])) + # 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.""" + 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): + """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(), + 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": + 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) + 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 + + +# 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", "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( + 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 new file mode 100644 index 000000000..9bdfd1343 --- /dev/null +++ b/tests/core/test_known_periods_branch_visibility.py @@ -0,0 +1,507 @@ +"""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``. +``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 + (``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, 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 +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, 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 + 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): + """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") + 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) +@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[0]) + 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 + + +@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") + 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_", +) +# 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"), + *( + pytest.param(periods.period(period), id=period, marks=skip_on_windows) + for period in ("month:2025-01:3", "year:2025-03", "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 + + +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) + ) + + +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 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) + }