From 692bbd6abb6856d8ba9e9105d71394649d9c2a08 Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Tue, 13 Oct 2026 06:34:31 -0400 Subject: [PATCH 1/6] Stamp saved US outputs with the rename register they were calculated under A US run records, before applying the legacy-input mapping, a stamp holding the whole LEGACY_INPUT_RENAMES register, an output format revision and the year-file stamp of its input dataset (null when it keeps none). save() writes it last as policyengine_output_stamp; load() refuses an output whose stamp is not exactly this release's, so Simulation.ensure() calculates it again. PolicyEngineUSDataset.load() and create_datasets() keep a year file's stamp in the dataset's metadata so a run can record it. Fixes #592. Co-Authored-By: Claude Opus 5.5 --- changelog.d/592.fixed.md | 1 + docs/microsim.md | 59 +++- .../common/model_version.py | 60 +++- .../tax_benefit_models/us/datasets.py | 43 ++- .../tax_benefit_models/us/legacy_inputs.py | 313 +++++++++++++++++- .../tax_benefit_models/us/model.py | 9 + tests/test_us_legacy_inputs_integration.py | 11 +- 7 files changed, 445 insertions(+), 51 deletions(-) create mode 100644 changelog.d/592.fixed.md diff --git a/changelog.d/592.fixed.md b/changelog.d/592.fixed.md new file mode 100644 index 00000000..204b811e --- /dev/null +++ b/changelog.d/592.fixed.md @@ -0,0 +1 @@ +A US `Simulation.run()` stamps its output with the register of renamed stored inputs it was calculated under, an output format revision, and the stamp of the year file its input data was cut into, if any. `Simulation.load()` refuses a saved US output, and `Simulation.ensure()` runs it again, unless that stamp matches this release exactly. An output calculated before a rename was registered, or from a year file that is out of date, is therefore not reused. The error says why, naming the renames the output may lack. Outputs saved before the stamp existed are calculated again once. diff --git a/docs/microsim.md b/docs/microsim.md index 1e63432a..97cdd6af 100644 --- a/docs/microsim.md +++ b/docs/microsim.md @@ -401,11 +401,9 @@ The renames applied are recorded as `{old: new}` (`{}` when none applied): carried the draw: data that stores neither name runs with the new input's default, which for WIC is full take-up. -A saved US output written before this mapping existed has no record and may -have been calculated without the mapped inputs, so `load()` refuses it, and -`ensure()` runs the simulation again and saves the new output. Saved outputs -are checked only for the record, not for the register they were calculated -under. +Year files and saved outputs are each stamped with the register they were +cut or calculated under, so neither is reused once the register changes (see +below). ### Year files @@ -443,7 +441,56 @@ than cutting the files again, when a year file cannot be opened for reasons unrelated to its contents: no permission, a directory at its path, or a lock held by another process. A year file opened directly, as `PolicyEngineUSDataset(filepath=...)`, is not checked, and saving it again with -`save()` drops its stamp. +`save()` drops its stamp. A dataset read from a year file, or returned by +`create_datasets`, keeps the file's stamp as `metadata["year_file_stamp"]`, so +an output calculated from it records the stamp (see below). + +### Saved outputs + +A US `Simulation.run()` records a stamp in its output's +`metadata["output_stamp"]`, read before it applies the mapping: UTF-8 JSON +`{"register": {old: new, ...}, "revision": 1, "year_file": ...}`. It holds +the whole register, an output format revision, and the stamp of the year file +the input data was cut into, or `null` when the input dataset keeps none (an +in-memory dataset or a source file, for example). `save()` writes it into the +output file, after everything else, as +`legacy_inputs.OUTPUT_STAMP_H5_DATASET`; `load()` restores it. Its name and +keys differ from the year-file stamp's, so an output never passes as a year +file, nor a year file as an output. + +`Simulation.load()` refuses a saved US output unless its stamp is exactly what +this release would write for a run over a dataset with no year-file stamp, or +over a current year file, and its renames record lists only renames in the +stamped register. `Simulation.ensure()` then runs the simulation again and +saves the new output. `legacy_inputs.output_stamp_problem(path)` returns the +reason, or `None` for a current output; it reads only the stamp and the +record. A saved output is out of date when: + +- it has no stamp. Outputs saved before the mapping existed have no record + either, and may have been calculated with every WIC-eligible person taking + WIC up. Outputs saved with the record but before the stamp are calculated + again once, and so is an output whose save did not finish; +- its register differs from this release's. An output calculated before an + entry was registered may lack that entry's input, and one calculated under + an entry this release lacks, or maps to another new name, used a mapping + this release does not; +- it was calculated from a year file whose stamp is not this release's: a + run over a year file maps only what the file still stores under an old + name, so the output inherits whatever the file lost. That happens for a + year file opened directly when `load_datasets` would refuse it, and for + every year file once the register or the year-file format revision + changes; +- it has another output format revision, or its stamp or record cannot be + read. + +So any change to `LEGACY_INPUT_RENAMES` makes every saved US output out of +date too. A missing output file raises `FileNotFoundError`, as before, so +`ensure()` runs the simulation. An output file that cannot be opened for +reasons unrelated to its contents (no permission, a directory, a lock held by +another process) raises that error rather than counting as out of date. A +year file cut before the stamp existed and opened directly keeps no stamp, so +an output calculated from it records `null` and is reused; `load_datasets` +refuses such a file. ## Pinned model versions diff --git a/src/policyengine/tax_benefit_models/common/model_version.py b/src/policyengine/tax_benefit_models/common/model_version.py index e10d3330..5cf6a581 100644 --- a/src/policyengine/tax_benefit_models/common/model_version.py +++ b/src/policyengine/tax_benefit_models/common/model_version.py @@ -356,9 +356,11 @@ def save(self, simulation: Simulation) -> None: "something to persist." ) serialized_spm = None + stamp = None if self.country_code == "us": from policyengine.core.spm import SPMProvenance from policyengine.tax_benefit_models.us.legacy_inputs import ( + OUTPUT_STAMP_KEY, RENAMES_RECORD_KEY, ) @@ -379,15 +381,23 @@ def save(self, simulation: Simulation) -> None: ) # ``PolicyEngineUSDataset.save()`` writes this record into the # file, and ``load()`` refuses an output without it. - renames = (getattr(simulation.output_dataset, "metadata", None) or {}).get( - RENAMES_RECORD_KEY - ) + metadata = getattr(simulation.output_dataset, "metadata", None) or {} + renames = metadata.get(RENAMES_RECORD_KEY) if renames is None: raise ValueError( "This US output does not record which renamed stored " "inputs were mapped when it was calculated; run again " "before saving" ) + # The stamp the run recorded, not one made now: the register the + # mapping used when the output was calculated. + stamp = metadata.get(OUTPUT_STAMP_KEY) + if stamp is None: + raise ValueError( + "This US output does not record the register of renamed " + "stored inputs it was calculated under; run again before " + "saving" + ) simulation.output_dataset.save() if serialized_spm is not None: # Store UTF-8 JSON in a dataset rather than an attribute: the @@ -399,6 +409,15 @@ def save(self, simulation: Simulation) -> None: data=serialized_spm, dtype=h5py.string_dtype("utf-8"), ) + if stamp is not None: + from policyengine.tax_benefit_models.us.legacy_inputs import ( + write_output_stamp, + ) + + # Stamp last. ``PolicyEngineUSDataset.save()`` truncates the file + # first, so an output whose save stopped part-way has no stamp, + # and ``load()`` refuses it. + write_output_stamp(simulation.output_dataset.filepath, stamp) def load(self, simulation: Simulation) -> None: """Rehydrate the simulation's output dataset from disk. @@ -414,17 +433,36 @@ def load(self, simulation: Simulation) -> None: if self.country_code == "us": from policyengine.core.spm import SPMProvenance from policyengine.tax_benefit_models.us.legacy_inputs import ( + OUTPUT_STAMP_KEY, RENAMES_RECORD_KEY, + output_stamp_problem, + read_output_stamp, read_renames_record, ) + # Outputs calculated under another register of renamed stored + # inputs, or from a year file cut under one, may have been + # calculated without a mapped input (for example with every + # WIC-eligible person taking WIC up), so they are not reused. + # ``Simulation.ensure()`` runs such a simulation again. A missing + # file raises ``FileNotFoundError`` here. + problem = output_stamp_problem(filepath) + if problem is not None: + raise ValueError( + f"Saved US simulation output {filepath} is out of date: " + f"{problem}. Run the simulation again (Simulation.ensure() " + "does so)." + ) with h5py.File(filepath, "r") as stream: raw = ( stream["policyengine_spm"].asstr()[()] if "policyengine_spm" in stream else None ) + # Both are current: ``output_stamp_problem`` accepts an output + # only with a stamp and a well-formed record. recorded_renames = read_renames_record(filepath) + recorded_stamp = read_output_stamp(filepath) if raw is None: raise ValueError( "Saved US simulation has no SPM configuration or receipt" @@ -433,15 +471,6 @@ def load(self, simulation: Simulation) -> None: if recorded["config"] != simulation.spm_config: raise ValueError("Saved US simulation uses different SPM settings") receipt = SPMProvenance.model_validate(recorded["provenance"]) - # Outputs saved before stored inputs were mapped onto renamed - # live inputs were calculated without them (for example with - # every WIC-eligible person taking WIC up), so they are not - # reused. ``Simulation.ensure()`` runs such a simulation again. - if recorded_renames is None: - raise ValueError( - "Saved US simulation predates the mapping of renamed " - "stored inputs (it has no record of them); run it again" - ) if self.country_code == "uk": from policyengine.tax_benefit_models.uk.datasets import ( @@ -471,8 +500,13 @@ def load(self, simulation: Simulation) -> None: simulation.spm_receipt = receipt simulation.spm = SPMSelection.model_validate(recorded["config"]) + # Kept so that saving the loaded output again writes its stamp. simulation.output_dataset.metadata.update( - {"spm_config": recorded["config"], RENAMES_RECORD_KEY: recorded_renames} + { + "spm_config": recorded["config"], + RENAMES_RECORD_KEY: recorded_renames, + OUTPUT_STAMP_KEY: recorded_stamp, + } ) if os.path.exists(filepath): diff --git a/src/policyengine/tax_benefit_models/us/datasets.py b/src/policyengine/tax_benefit_models/us/datasets.py index a5fbfb9b..4233747a 100644 --- a/src/policyengine/tax_benefit_models/us/datasets.py +++ b/src/policyengine/tax_benefit_models/us/datasets.py @@ -1,3 +1,4 @@ +import copy import hashlib import importlib.util import json @@ -27,6 +28,7 @@ ) from policyengine.tax_benefit_models.us.legacy_inputs import ( RENAMES_RECORD_KEY, + YEAR_FILE_STAMP_KEY, apply_legacy_input_renames_to_microsimulation, check_yearly_periods, pending_legacy_input_renames, @@ -34,6 +36,7 @@ write_renames_record, write_year_file_stamp, year_file_stamp, + year_file_stamp_metadata, year_file_stamp_problem, ) from policyengine.utils.hashing import sha256_file @@ -127,20 +130,29 @@ def load(self) -> None: self.metadata[RENAMES_RECORD_KEY] = renames if _is_policyengine_core_h5(path): self.data = _load_policyengine_core_h5(path, self.year) - return - - with pd.HDFStore(filepath, mode="r") as store: - data = {entity: store[entity] for entity in US_ENTITY_KEYS} - # Native producer files store calibrated household weights. Project - # missing weights through membership IDs on these in-memory copies, - # before constructing MicroDataFrames; the source HDF remains read-only. - _assign_missing_entity_weights(data) - self.data = USYearData( - **{ - entity: MicroDataFrame(frame, weights=f"{entity}_weight") - for entity, frame in data.items() - } - ) + else: + with pd.HDFStore(filepath, mode="r") as store: + data = {entity: store[entity] for entity in US_ENTITY_KEYS} + # Native producer files store calibrated household weights. Project + # missing weights through membership IDs on these in-memory copies, + # before constructing MicroDataFrames; the source HDF remains + # read-only. + _assign_missing_entity_weights(data) + self.data = USYearData( + **{ + entity: MicroDataFrame(frame, weights=f"{entity}_weight") + for entity, frame in data.items() + } + ) + # The stamp of the year file this data was cut into, if it is one. A + # run records it in its output's stamp, so an output calculated from + # a year file that goes out of date is not reused (see + # ``legacy_inputs.OUTPUT_STAMP_H5_DATASET``). + stamp = year_file_stamp_metadata(path) + if stamp is None: + self.metadata.pop(YEAR_FILE_STAMP_KEY, None) + else: + self.metadata[YEAR_FILE_STAMP_KEY] = stamp def __repr__(self) -> str: if self.data is None: @@ -572,6 +584,9 @@ def create_datasets( # Stamp last. ``save()`` truncates the file first, so a year file # whose write stopped part-way has no stamp and is cut again. write_year_file_stamp(us_dataset.filepath, stamp) + # Keep it, as loading the file would, so a run over the returned + # dataset records it in its output's stamp. + us_dataset.metadata[YEAR_FILE_STAMP_KEY] = copy.deepcopy(stamp) dataset_key = f"{dataset_stem}_{year}" result[dataset_key] = us_dataset diff --git a/src/policyengine/tax_benefit_models/us/legacy_inputs.py b/src/policyengine/tax_benefit_models/us/legacy_inputs.py index 38fb0249..5607ddf3 100644 --- a/src/policyengine/tax_benefit_models/us/legacy_inputs.py +++ b/src/policyengine/tax_benefit_models/us/legacy_inputs.py @@ -22,6 +22,7 @@ from __future__ import annotations +import copy import inspect import json from collections.abc import Iterable, Iterator, Mapping @@ -41,11 +42,13 @@ #: ``would_claim_wic`` only, so without this mapping every WIC-eligible person #: takes WIC up (PolicyEngine/microcosm#1026). #: -#: Each ``create_datasets`` year file is stamped with this register (see -#: :data:`YEAR_FILE_STAMP_H5_DATASET`), so any change to it (an entry added, -#: removed or pointed at another live input) makes every year file stale, and -#: ``ensure_datasets`` cuts them all again. An entry that no longer applies -#: turns itself off, so leave it in place rather than remove it. +#: Each ``create_datasets`` year file and each saved output is stamped with +#: this register (see :data:`YEAR_FILE_STAMP_H5_DATASET` and +#: :data:`OUTPUT_STAMP_H5_DATASET`), so any change to it (an entry added, +#: removed or pointed at another live input) makes every year file and every +#: saved output stale: ``ensure_datasets`` cuts the year files again, and +#: ``Simulation.ensure()`` calculates the outputs again. An entry that no +#: longer applies turns itself off, so leave it in place rather than remove it. LEGACY_INPUT_RENAMES: dict[str, str] = { "would_claim_wic": "takes_up_wic_if_eligible", } @@ -68,10 +71,9 @@ #: ``ensure_datasets`` regenerates it. #: #: The record shows that a file was written with the mapping, not under which -#: register. Year files also carry :data:`YEAR_FILE_STAMP_H5_DATASET`, which -#: ties them to the register they were cut under. Saved outputs carry only -#: this record, so an output saved before a register entry was added is still -#: reused. +#: register. Year files also carry :data:`YEAR_FILE_STAMP_H5_DATASET`, and +#: saved outputs :data:`OUTPUT_STAMP_H5_DATASET`, which tie each to the +#: register it was cut or calculated under. RENAMES_H5_DATASET = "policyengine_legacy_input_renames" #: H5 dataset stamping a ``create_datasets`` year file, as UTF-8 JSON @@ -94,8 +96,48 @@ #: Removing the temporary ACS WIC assumption (:func:`_temporary_acs_wic_values`) #: is such a change: year files cut while it applies store assumed claiming #: as ``takes_up_wic_if_eligible`` for ACS people with no stored decision. +#: +#: A saved output records the stamp of the year file it was calculated from, +#: so raising this makes those outputs out of date too. YEAR_FILE_STAMP_REVISION = 1 +#: Key under which a dataset read from, or cut into, a year file keeps that +#: file's stamp in its ``metadata``, as the JSON value stored (see +#: :func:`read_year_file_stamp`). A run records it in its output's stamp. +YEAR_FILE_STAMP_KEY = "year_file_stamp" + +#: H5 dataset stamping a saved US simulation output, as UTF-8 JSON +#: ``{"register": {legacy: live, ...}, "revision": int, "year_file": ...}``: +#: the whole register the output was calculated under, its output format +#: revision, and the stamp of the year file it was calculated from (``null`` +#: when the input dataset carried none). See :func:`output_stamp`. +#: +#: A US run records the stamp in its output's ``metadata`` (under +#: :data:`OUTPUT_STAMP_KEY`) when it is calculated, before applying the +#: mapping. ``MicrosimulationModelVersion.save()`` writes it after the tables, +#: the renames record and the SPM receipt, so an output whose save stopped +#: part-way has none. ``Simulation.load()`` refuses an output whose stamp does +#: not match this release's (see :func:`output_stamp_problem`), and +#: ``Simulation.ensure()`` then calculates it again. +#: +#: Its name and its keys differ from :data:`YEAR_FILE_STAMP_H5_DATASET`'s, so +#: an output never passes as a year file, nor a year file as an output. +OUTPUT_STAMP_H5_DATASET = "policyengine_output_stamp" + +#: Format revision of saved outputs, stored in their stamp. Raise it when a +#: change to the mapping, or to how a run uses it, means that outputs +#: calculated before it must be calculated again, for a reason neither the +#: register nor the year-file stamp captures. +#: +#: Removing the temporary ACS WIC assumption (:func:`_temporary_acs_wic_values`) +#: is such a change: a run over a source that stores ``would_claim_wic`` +#: assumes claiming for ACS people with no stored decision while it applies. +#: Raise :data:`YEAR_FILE_STAMP_REVISION` with it. +OUTPUT_STAMP_REVISION = 1 + +#: Key under which a US run's output keeps its stamp in ``metadata``. +OUTPUT_STAMP_KEY = "output_stamp" + #: ``{year: {entity key: stored table}}``, one table per entity per dataset #: year, in the dataset's own row order. StoredTables = Mapping[int, Mapping[str, pd.DataFrame]] @@ -270,12 +312,7 @@ def year_file_stamp_problem(path: str | Path) -> str | None: ) if stamp == current: return None - missing = { - legacy: live for legacy, live in register.items() if stamped.get(legacy) != live - } - unknown = { - legacy: live for legacy, live in stamped.items() if register.get(legacy) != live - } + missing, unknown = _register_differences(stamped, register) reasons = [] if missing: reasons.append( @@ -290,6 +327,252 @@ def year_file_stamp_problem(path: str | Path) -> str | None: return "; ".join(reasons) +def year_file_stamp_metadata(path: str | Path) -> Any | None: + """Return the year-file stamp a dataset read from ``path`` keeps. + + A dataset keeps it in ``metadata`` under :data:`YEAR_FILE_STAMP_KEY`, and + a run records it in its output's stamp (see :func:`output_stamp`). + + An error that says nothing about the file's contents (a missing file, no + permission, a directory, a lock) is raised, as by + :func:`year_file_stamp_problem`. + + Returns: + ``None`` when the file stores no stamp or is not an H5 file; the + stored JSON value when it is readable and not ``null``; otherwise + ``"unreadable"``. No release stamps a year file with a string, so an + output calculated from a file whose stamp is unreadable is never + current. + """ + import h5py + + try: + with h5py.File(path, "r") as file: + stamp = _read_json_dataset(file, YEAR_FILE_STAMP_H5_DATASET) + except (FileNotFoundError, PermissionError, IsADirectoryError, BlockingIOError): + raise + except OSError: + return None + if stamp is _ABSENT: + return None + return "unreadable" if stamp is _UNREADABLE or stamp is None else stamp + + +def output_stamp(year_file: Any = None) -> dict[str, Any]: + """Return the stamp for an output calculated now. + + ``year_file`` is the stamp the input dataset keeps under + :data:`YEAR_FILE_STAMP_KEY` (see :func:`year_file_stamp_metadata`), or + ``None`` when it keeps none. Like :func:`year_file_stamp`, the stamp holds + the whole register, read when this is called. + """ + return { + "register": dict(sorted(LEGACY_INPUT_RENAMES.items())), + "revision": OUTPUT_STAMP_REVISION, + "year_file": copy.deepcopy(year_file), + } + + +def read_output_stamp(path: str | Path) -> Any | None: + """Return the JSON value of a saved output's stamp, well formed or not. + + ``None`` means the file stores no readable stamp (see + :data:`OUTPUT_STAMP_H5_DATASET`) or is not an H5 file. + """ + import h5py + + try: + with h5py.File(path, "r") as file: + stamp = _read_json_dataset(file, OUTPUT_STAMP_H5_DATASET) + except OSError: + return None + return None if stamp is _ABSENT or stamp is _UNREADABLE else stamp + + +def write_output_stamp(path: str | Path, stamp: Mapping[str, Any]) -> None: + """Stamp a saved output (see :func:`output_stamp`), replacing any stamp.""" + _write_json_dataset(path, OUTPUT_STAMP_H5_DATASET, dict(stamp)) + + +def output_stamp_problem(path: str | Path) -> str | None: + """Return why a saved US simulation output is out of date, or ``None``. + + A saved output is current only when both of these hold: + + - its stamp equals :func:`output_stamp` for an input that kept no + year-file stamp, or for one cut into a year file under this release's + :func:`year_file_stamp`. That is, it has the same output format + revision and the same register, and the stamp of the year file it was + calculated from is ``null`` or this release's; + - it stores the renames record, listing only renames in its stamped + register. + + The registers must match exactly, for the reasons + :func:`year_file_stamp_problem` gives. An output calculated before an + entry was registered may lack that entry's input, as outputs calculated + before the WIC entry gave every WIC-eligible person WIC. A run over a year + file maps only what the file still stores under a legacy name, so an + output calculated from a year file that is out of date under this release + is out of date too. + + Only the stamp and the record are read, never the tables. A file that + cannot be read as HDF5 is out of date. A missing file raises + ``FileNotFoundError``, so ``Simulation.ensure()`` calculates the output + without reporting an unexpected error. An error that says nothing about + the file's contents (no permission, a directory, a file another process + has locked) is raised too, since the output could not be saved again + either. + + Returns: + ``None`` for a current output; otherwise a clause, starting "it" or + "its", that says why the output is out of date and names the renames + it may lack. + """ + import h5py + + register = dict(sorted(LEGACY_INPUT_RENAMES.items())) + current_year_file = year_file_stamp() + try: + with h5py.File(path, "r") as file: + stamp = _read_json_dataset(file, OUTPUT_STAMP_H5_DATASET) + record = _read_json_dataset(file, RENAMES_H5_DATASET) + except (FileNotFoundError, PermissionError, IsADirectoryError, BlockingIOError): + raise + except OSError: + return "it cannot be read as an HDF5 file (it may be half-written or corrupt)" + + # The revision comes first: a later revision may shape its stamp and + # record differently. + revision = stamp.get("revision") if isinstance(stamp, dict) else None + if type(revision) is int and revision != OUTPUT_STAMP_REVISION: + return ( + f"it was saved with output format revision {revision}, and this " + f"version of policyengine.py reads only revision " + f"{OUTPUT_STAMP_REVISION}" + ) + if stamp is _ABSENT: + if record is _ABSENT: + return ( + "it has no record of the renamed stored inputs mapped when it " + "was calculated, so it was saved by an older policyengine.py " + "or not saved completely" + + ( + f", and may have been calculated without any of: " + f"{_describe(register)}" + if register + else "" + ) + ) + return ( + "it has no stamp of the register of renamed inputs it was " + "calculated under" + + ( + f", so it may have been calculated without any of: " + f"{_describe(register)}" + if register + else "" + ) + ) + # ``type(...) is int`` rather than ``isinstance``: JSON ``true`` loads as + # ``True``, which equals 1. + if not ( + isinstance(stamp, dict) + and set(stamp) == {"register", "revision", "year_file"} + and type(revision) is int + and _is_str_map(stamp["register"]) + ): + return ( + "its stamp of the register of renamed inputs it was calculated " + "under is malformed" + ) + if record is _ABSENT: + return ( + "it has no record of the renamed stored inputs mapped when it was " + "calculated" + ) + if not _is_str_map(record): + return ( + "its record of the renamed stored inputs mapped when it was " + "calculated is malformed" + ) + stamped = stamp["register"] + outside = dict(record.items() - stamped.items()) + if outside: + return ( + "its record lists renames outside the register it was stamped " + f"with: {_describe(outside)}" + ) + missing, unknown = _register_differences(stamped, register) + reasons = [] + if missing: + reasons.append( + "it was calculated before these renames were registered, so it " + f"may have been calculated without them: {_describe(missing)}" + ) + if unknown: + reasons.append( + "it was calculated under renames this version of policyengine.py " + f"does not register: {_describe(unknown)}" + ) + if stamp["year_file"] is not None: + reasons.extend(_year_file_reasons(stamp["year_file"], current_year_file)) + return "; ".join(reasons) or None + + +def _year_file_reasons(stamp: Any, current: Mapping[str, Any]) -> list[str]: + """Say why the stamp of an output's year file is not ``current``. + + Returns no reason exactly when ``stamp`` equals ``current``. + """ + revision = stamp.get("revision") if isinstance(stamp, dict) else None + if type(revision) is int and revision != current["revision"]: + return [ + "it was calculated from a year file cut with year-file format " + f"revision {revision}, and this version of policyengine.py reads " + f"only revision {current['revision']}" + ] + if not ( + isinstance(stamp, dict) + and set(stamp) == {"register", "revision"} + and type(revision) is int + and _is_str_map(stamp["register"]) + ): + return [ + "its record of the stamp of the year file it was calculated from " + "is malformed" + ] + missing, unknown = _register_differences(stamp["register"], current["register"]) + reasons = [] + if missing: + reasons.append( + "it was calculated from a year file cut before these renames were " + f"registered, which may have lost them: {_describe(missing)}" + ) + if unknown: + reasons.append( + "it was calculated from a year file cut under renames this version " + f"of policyengine.py does not register: {_describe(unknown)}" + ) + return reasons + + +def _register_differences( + stamped: Mapping[str, str], register: Mapping[str, str] +) -> tuple[dict[str, str], dict[str, str]]: + """Return the entries each register has that the other lacks or retargets. + + The first holds ``register``'s entries that ``stamped`` lacks or maps to + another live input; the second, ``stamped``'s that ``register`` does. + """ + missing = { + legacy: live for legacy, live in register.items() if stamped.get(legacy) != live + } + unknown = { + legacy: live for legacy, live in stamped.items() if register.get(legacy) != live + } + return missing, unknown + + def _read_json_dataset(file: Any, name: str) -> Any: """Return the JSON value stored in an open H5 file's ``name`` dataset.""" import h5py diff --git a/src/policyengine/tax_benefit_models/us/model.py b/src/policyengine/tax_benefit_models/us/model.py index f07f1611..5246e58b 100644 --- a/src/policyengine/tax_benefit_models/us/model.py +++ b/src/policyengine/tax_benefit_models/us/model.py @@ -18,9 +18,12 @@ from .datasets import PolicyEngineUSDataset, USYearData, _validate_entity_ids from .legacy_inputs import ( + OUTPUT_STAMP_KEY, RENAMES_RECORD_KEY, + YEAR_FILE_STAMP_KEY, apply_legacy_input_renames, apply_legacy_input_renames_to_microsimulation, + output_stamp, ) from .spm import ( SPMProvenance, @@ -224,6 +227,11 @@ class InputMicrosimulation(Microsimulation): # leaves the module-level one untouched. Building populations # against the module-level system would hide reform-registered # variables like ``ctc_minimum_refundable_amount`` at calc time. + # The register the mapping below uses, and the stamp of the year file + # the input data was cut into, if it is one, stamped into the output + # when it is saved. ``load()`` refuses the output once either is out + # of date (see ``legacy_inputs.OUTPUT_STAMP_H5_DATASET``). + stamp = output_stamp(simulation.dataset.metadata.get(YEAR_FILE_STAMP_KEY)) # Renames already applied when the input data was cut (a # ``create_datasets`` year file stores the mapped input under its # live name) count as applied to this run's inputs too, so the @@ -335,6 +343,7 @@ class InputMicrosimulation(Microsimulation): metadata={ "spm_config": dict(microsim.spm_config), RENAMES_RECORD_KEY: dict(sorted(legacy_input_renames.items())), + OUTPUT_STAMP_KEY: stamp, }, data=USYearData( person=data["person"], diff --git a/tests/test_us_legacy_inputs_integration.py b/tests/test_us_legacy_inputs_integration.py index 182d39af..040c379e 100644 --- a/tests/test_us_legacy_inputs_integration.py +++ b/tests/test_us_legacy_inputs_integration.py @@ -872,19 +872,24 @@ def test_an_output_saved_before_the_mapping_is_not_reused(tmp_path): ``load()`` refuses it, and ``ensure()`` runs the simulation again. """ from policyengine.core.simulation import _cache - from policyengine.tax_benefit_models.us.legacy_inputs import RENAMES_H5_DATASET + from policyengine.tax_benefit_models.us.legacy_inputs import ( + OUTPUT_STAMP_H5_DATASET, + RENAMES_H5_DATASET, + ) simulation = _saved_run(tmp_path, "legacy-wic-stale") path = simulation.output_dataset.filepath # Make the file look like one saved before the mapping existed: no - # record, and the full-take-up WIC every eligible person got then. + # record or stamp, and the full-take-up WIC every eligible person got + # then. with h5py.File(path, "a") as stream: del stream[RENAMES_H5_DATASET] + del stream[OUTPUT_STAMP_H5_DATASET] stale = pd.DataFrame(pd.read_hdf(path, "person")) stale["wic"] = stale["wic"].where(stale["person_id"] != 2, 999.0) stale.to_hdf(path, key="person", format="fixed") - with pytest.raises(ValueError, match="predates the mapping"): + with pytest.raises(ValueError, match="saved by an older policyengine.py"): _reloaded(simulation) _cache.clear() From 66e97bb4a35f33ffacc277d77764d92dfaec362e Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Tue, 13 Oct 2026 06:46:33 -0400 Subject: [PATCH 2/6] Test the saved-output stamp: properties, real-engine runs, load/ensure agreement Engine-free unit and Hypothesis tests state exactness, totality, round trip, named renames, truncation and stamp distinctness for output_stamp_problem. Integration tests over policyengine-us 2.2.1 with the three-person WIC fixture replay the WIC history for a saved output, cover year-file provenance, save ordering and partial saves, and check that Simulation.load() refuses, and ensure() runs again, exactly when output_stamp_problem gives a reason. Co-Authored-By: Claude Opus 5.5 --- tests/test_us_legacy_inputs_integration.py | 470 ++++++++- tests/test_us_output_stamp.py | 1079 ++++++++++++++++++++ 2 files changed, 1548 insertions(+), 1 deletion(-) create mode 100644 tests/test_us_output_stamp.py diff --git a/tests/test_us_legacy_inputs_integration.py b/tests/test_us_legacy_inputs_integration.py index 040c379e..0ff941ab 100644 --- a/tests/test_us_legacy_inputs_integration.py +++ b/tests/test_us_legacy_inputs_integration.py @@ -13,13 +13,17 @@ from __future__ import annotations +import json +import shutil +import tempfile +from pathlib import Path from unittest.mock import patch import h5py import numpy as np import pandas as pd import pytest -from hypothesis import given, settings +from hypothesis import event, given, settings from hypothesis import strategies as st from microdf import MicroDataFrame @@ -37,10 +41,14 @@ load_datasets, ) from policyengine.tax_benefit_models.us.legacy_inputs import ( # noqa: E402 + OUTPUT_STAMP_H5_DATASET, RENAMES_H5_DATASET, YEAR_FILE_STAMP_H5_DATASET, apply_legacy_input_renames, apply_legacy_input_renames_to_microsimulation, + output_stamp, + output_stamp_problem, + read_output_stamp, read_renames_record, read_year_file_stamp, year_file_stamp, @@ -927,3 +935,463 @@ def test_the_run_record_binds_the_renames_applied(tmp_path): payloads = build_simulation_run_record_payloads(simulation) assert payloads["results"]["legacy_input_renames"] == RENAME + + +# --- Saved outputs are stamped with the register ------------------------ + +FUTURE = {"fictitious_old_input": "fictitious_new_input"} +OUT_OF_DATE = "Saved US simulation output {path} is out of date: {reason}. Run the simulation again (Simulation.ensure() does so)." + + +def _stamped_run(dataset, simulation_id): + """Run over ``dataset`` and save, as ``ensure()`` does on a miss.""" + simulation = pe.Simulation( + id=simulation_id, + dataset=dataset, + tax_benefit_model_version=pe.us.model, + extra_variables={"person": [LIVE, "wic"]}, + ) + simulation.run() + simulation.save() + return simulation + + +def _ensure_output(simulation, dataset=None): + """``ensure()`` a fresh copy of ``simulation``, past the in-process cache. + + Returns: + The ensured simulation, and whether ``ensure()`` ran it again. + """ + from policyengine.core.simulation import _cache + + runs = [] + run = pe.Simulation.run + + def counting(self): + runs.append(self.id) + return run(self) + + _cache.clear() + try: + ensured = pe.Simulation( + id=simulation.id, + dataset=dataset if dataset is not None else simulation.dataset, + tax_benefit_model_version=pe.us.model, + extra_variables=simulation.extra_variables, + ) + with patch.object(pe.Simulation, "run", counting): + ensured.ensure() + finally: + _cache.clear() + return ensured, bool(runs) + + +def test_a_saved_output_is_stamped_with_the_register_it_was_calculated_under( + tmp_path, +): + simulation = _saved_run(tmp_path, "stamped-output") + path = simulation.output_dataset.filepath + + # An in-memory input keeps no year-file stamp. + assert simulation.output_dataset.metadata["output_stamp"] == output_stamp() + assert output_stamp() == {"register": RENAME, "revision": 1, "year_file": None} + assert read_output_stamp(path) == output_stamp() + assert output_stamp_problem(path) is None + + restored = _reloaded(simulation) + assert restored.output_dataset.metadata["output_stamp"] == output_stamp() + # Saving the loaded output again keeps its stamp, so it loads again. + restored.save() + assert read_output_stamp(path) == output_stamp() + assert _person_outputs(_reloaded(restored))["wic"].iloc[INFANT] == 0 + _, ran = _ensure_output(simulation) + assert not ran + + +def test_an_output_calculated_before_a_register_entry_is_calculated_again(tmp_path): + """Replay the WIC history for a saved output that has the record. + + Calculating under an empty register is what policyengine.py did before + the WIC entry existed. The record check alone accepted such an output, + since it records ``{}``. + """ + with patch.object(legacy_inputs, "LEGACY_INPUT_RENAMES", {}): + simulation = _saved_run(tmp_path, "before-the-entry") + path = simulation.output_dataset.filepath + # Under the register it was calculated with, the output is current. + assert output_stamp_problem(path) is None + # Without the mapping, the infant takes WIC up. + assert _person_outputs(simulation)["wic"].iloc[INFANT] > 0 + assert read_renames_record(path) == {} + assert read_output_stamp(path) == {"register": {}, "revision": 1, "year_file": None} + + reason = ( + "it was calculated before these renames were registered, so it may have " + f"been calculated without them: {LEGACY} -> {LIVE}" + ) + assert output_stamp_problem(path) == reason + with pytest.raises(ValueError) as refused: + _reloaded(simulation) + assert str(refused.value) == OUT_OF_DATE.format(path=path, reason=reason) + + ensured, ran = _ensure_output(simulation) + + assert ran + assert _person_outputs(ensured)["wic"].iloc[INFANT] == 0 + assert read_output_stamp(path) == output_stamp() + # The new output is current, so it is reused from now on. + reused, ran = _ensure_output(simulation) + assert not ran + assert _person_outputs(reused)["wic"].iloc[INFANT] == 0 + + +def test_a_future_register_entry_makes_saved_outputs_stale(tmp_path): + """An entry the engine does not need still makes older outputs stale.""" + from policyengine_us.system import system + + for name in (*FUTURE, *FUTURE.values()): + assert name not in system.variables + simulation = _saved_run(tmp_path, "future-entry") + path = simulation.output_dataset.filepath + + with patch.object(legacy_inputs, "LEGACY_INPUT_RENAMES", {**RENAME, **FUTURE}): + assert output_stamp_problem(path) == ( + "it was calculated before these renames were registered, so it may " + "have been calculated without them: fictitious_old_input -> " + "fictitious_new_input" + ) + with pytest.raises(ValueError, match="fictitious_old_input"): + _reloaded(simulation) + + ensured, ran = _ensure_output(simulation) + + assert ran + assert read_output_stamp(path) == { + "register": {**RENAME, **FUTURE}, + "revision": 1, + "year_file": None, + } + # Only the entry the engine needs was applied, and the results are as + # before. + assert read_renames_record(path) == RENAME + assert _person_outputs(ensured)["wic"].iloc[INFANT] == 0 + _, ran = _ensure_output(simulation) + assert not ran + + # Back under this release's register, the newer output is out of date too. + assert "does not register: fictitious_old_input" in output_stamp_problem(path) + + +def test_an_output_saved_with_only_the_record_is_calculated_again(tmp_path): + """An output saved with the record but before the stamp existed.""" + simulation = _saved_run(tmp_path, "record-only") + path = simulation.output_dataset.filepath + with h5py.File(path, "a") as stream: + del stream[OUTPUT_STAMP_H5_DATASET] + + with pytest.raises(ValueError, match="it has no stamp of the register"): + _reloaded(simulation) + _, ran = _ensure_output(simulation) + + assert ran + assert output_stamp_problem(path) is None + + +def test_a_run_over_a_year_file_records_its_stamp(tmp_path): + source, created = _create_year_file(tmp_path) + # ``create_datasets`` and ``load_datasets`` (through ``ensure_datasets``) + # both keep the year file's stamp. + loaded = _assert_reused(source, tmp_path / "data") + for dataset in (created, loaded): + assert dataset.metadata["year_file_stamp"] == year_file_stamp() + + for index, dataset in enumerate((created, loaded)): + simulation = _stamped_run(dataset, f"from-a-year-file-{index}") + path = simulation.output_dataset.filepath + + stamp = output_stamp(year_file_stamp()) + assert stamp["year_file"] == {"register": RENAME, "revision": 1} + assert simulation.output_dataset.metadata["output_stamp"] == stamp + assert read_output_stamp(path) == stamp + assert output_stamp_problem(path) is None + assert _person_outputs(_reloaded(simulation))["wic"].iloc[INFANT] == 0 + + +def test_an_output_from_an_out_of_date_year_file_is_not_reused(tmp_path): + """A year file cut before the WIC entry lost the draw, so an output + calculated from it, under this release's register, is wrong too.""" + with patch.object(legacy_inputs, "LEGACY_INPUT_RENAMES", {}): + source, dataset = _create_year_file(tmp_path) + # Opened directly: ``load_datasets`` would refuse it. + stale = PolicyEngineUSDataset( + name="stale", description="stale", filepath=dataset.filepath, year=YEAR + ) + cut_under = {"register": {}, "revision": 1} + assert stale.metadata["year_file_stamp"] == cut_under + + simulation = _stamped_run(stale, "from-a-stale-year-file") + path = simulation.output_dataset.filepath + + # The file stores neither name, so there was nothing to map. + assert _person_outputs(simulation)["wic"].iloc[INFANT] > 0 + assert read_output_stamp(path) == output_stamp(cut_under) + # The output's own register is current; its year file's is not. + reason = ( + "it was calculated from a year file cut before these renames were " + f"registered, which may have lost them: {LEGACY} -> {LIVE}" + ) + assert output_stamp_problem(path) == reason + with pytest.raises(ValueError) as refused: + _reloaded(simulation) + assert str(refused.value) == OUT_OF_DATE.format(path=path, reason=reason) + + # Once ensure_datasets cuts the year file again, ensure() over it + # calculates the output again rather than reusing the stale one. + regenerated = _ensure(source, tmp_path / "data") + assert regenerated.filepath == dataset.filepath + ensured, ran = _ensure_output(simulation, dataset=regenerated) + + assert ran + assert _person_outputs(ensured)["wic"].iloc[INFANT] == 0 + assert read_output_stamp(path) == output_stamp(year_file_stamp()) + _, ran = _ensure_output(simulation, dataset=regenerated) + assert not ran + + +def test_save_writes_the_stamp_the_run_recorded(tmp_path): + """Not one made at save time, under whatever register is current then.""" + with patch.object(legacy_inputs, "LEGACY_INPUT_RENAMES", {}): + simulation = _run(_in_memory_dataset(tmp_path)) + simulation.save() + path = simulation.output_dataset.filepath + + assert read_output_stamp(path) == {"register": {}, "revision": 1, "year_file": None} + assert output_stamp_problem(path).startswith( + "it was calculated before these renames were registered" + ) + + +def test_an_output_without_a_stamp_is_not_saved(tmp_path): + simulation = _run(_in_memory_dataset(tmp_path)) + del simulation.output_dataset.metadata["output_stamp"] + + with pytest.raises( + ValueError, + match="register of renamed stored inputs it was calculated under; run " + "again before saving", + ): + simulation.save() + assert not Path(simulation.output_dataset.filepath).exists() + + +def test_save_stamps_the_output_last(tmp_path): + events = [] + save = PolicyEngineUSDataset.save + stamp = legacy_inputs.write_output_stamp + + def recording_save(self): + save(self) + events.append("tables") + + def recording_stamp(path, value): + with h5py.File(path, "r") as stream: + events.append( + ("stamp", RENAMES_H5_DATASET in stream, "policyengine_spm" in stream) + ) + stamp(path, value) + + with ( + patch.object(PolicyEngineUSDataset, "save", recording_save), + patch.object(legacy_inputs, "write_output_stamp", recording_stamp), + ): + simulation = _saved_run(tmp_path, "stamped-last") + + # The tables, the record and the SPM receipt are all written first. + assert events == ["tables", ("stamp", True, True)] + assert output_stamp_problem(simulation.output_dataset.filepath) is None + + +def test_an_output_whose_save_stopped_before_the_stamp_is_calculated_again( + tmp_path, +): + simulation = pe.Simulation( + id="stopped-before-stamping", + dataset=_in_memory_dataset(tmp_path), + tax_benefit_model_version=pe.us.model, + extra_variables={"person": [LIVE, "wic"]}, + ) + simulation.run() + + def stop(path, value): + raise RuntimeError("stopped before stamping") + + with patch.object(legacy_inputs, "write_output_stamp", stop): + with pytest.raises(RuntimeError, match="stopped before stamping"): + simulation.save() + + path = simulation.output_dataset.filepath + # Everything else was written, but without the stamp the output is + # calculated again. + assert read_renames_record(path) == RENAME + with h5py.File(path, "r") as stream: + assert "policyengine_spm" in stream + assert output_stamp_problem(path).startswith("it has no stamp") + _, ran = _ensure_output(simulation) + assert ran + assert output_stamp_problem(path) is None + + +def test_an_output_is_never_loaded_as_a_year_file_nor_the_reverse(tmp_path): + source, dataset = _create_year_file(tmp_path) + data_folder = tmp_path / "data" + simulation = _stamped_run(dataset, "beside-the-year-files") + + # A year file where an output is looked for is refused. + other = pe.Simulation( + id="a-year-file", + dataset=dataset, + tax_benefit_model_version=pe.us.model, + extra_variables={"person": [LIVE, "wic"]}, + ) + shutil.copy(dataset.filepath, data_folder / f"{other.storage_id}.h5") + with pytest.raises( + ValueError, + match="it has no stamp of the register of renamed inputs it was " + "calculated under", + ): + other.load() + + # An output where a year file is looked for is refused. + shutil.copy(simulation.output_dataset.filepath, dataset.filepath) + with pytest.raises( + ValueError, + match="it has no stamp of the register of renamed inputs it was cut under", + ): + load_datasets(datasets=[source], years=[YEAR], data_folder=str(data_folder)) + + +# --- Simulation.load() and ensure() agree with the stamp check ---------- + +#: Output stamps to store, around this release's. +STORED_STAMPS = [ + "absent", + {"register": RENAME, "revision": 1, "year_file": None}, + { + "register": RENAME, + "revision": 1, + "year_file": {"register": RENAME, "revision": 1}, + }, + {"register": {}, "revision": 1, "year_file": None}, + {"register": {**RENAME, **FUTURE}, "revision": 1, "year_file": None}, + {"register": RENAME, "revision": 1, "year_file": {"register": {}, "revision": 1}}, + { + "register": RENAME, + "revision": 1, + "year_file": {"register": RENAME, "revision": 2}, + }, + {"register": RENAME, "revision": 1, "year_file": "unreadable"}, + {"register": RENAME, "revision": 2, "year_file": None}, + {"register": RENAME, "revision": True, "year_file": None}, + {"register": RENAME, "revision": 1}, + "not json", +] +STORED_RECORDS = ["absent", RENAME, {}, {"x": "y"}, [], "not json"] +REGISTERS = [RENAME, {}, {**RENAME, **FUTURE}] + + +@pytest.fixture(scope="module") +def saved_output(tmp_path_factory): + """The bytes of a real saved output, and its file name.""" + simulation = _saved_run(tmp_path_factory.mktemp("saved_output"), "agreement") + path = Path(simulation.output_dataset.filepath) + return path.read_bytes(), path.name + + +def _restore(path, name, content): + with h5py.File(path, "a") as stream: + if name in stream: + del stream[name] + if content == "absent": + return + stream.create_dataset( + name, + data=content if content == "not json" else json.dumps(content), + dtype=h5py.string_dtype("utf-8"), + ) + + +@st.composite +def stored_outputs(draw): + """A register, and an output's stamp and record: current about a third of + the time, otherwise out of date in one of many ways.""" + register = draw(st.sampled_from(REGISTERS)) + current = [ + {"register": register, "revision": 1, "year_file": None}, + { + "register": register, + "revision": 1, + "year_file": {"register": register, "revision": 1}, + }, + ] + stamp = draw(st.sampled_from(current) | st.sampled_from(STORED_STAMPS)) + record = draw( + st.sampled_from([dict(register), {}]) | st.sampled_from(STORED_RECORDS) + ) + return register, stamp, record + + +@settings(max_examples=60, deadline=None) +@given(case=stored_outputs()) +def test_load_refuses_and_ensure_runs_again_exactly_when_an_output_is_out_of_date( + saved_output, case +): + """``Simulation.load()`` refuses a saved output exactly when + ``output_stamp_problem`` gives a reason, and says it; ``ensure()`` runs + the simulation again exactly then, and otherwise reuses the output.""" + from policyengine.core.simulation import _cache + + register, stamp, record = case + content, name = saved_output + with tempfile.TemporaryDirectory() as directory: + dataset = _in_memory_dataset(Path(directory)) + + def simulation(): + return pe.Simulation( + id="agreement", + dataset=dataset, + tax_benefit_model_version=pe.us.model, + extra_variables={"person": [LIVE, "wic"]}, + ) + + path = Path(directory) / name + assert path.name == f"{simulation().storage_id}.h5" + path.write_bytes(content) + _restore(path, OUTPUT_STAMP_H5_DATASET, stamp) + _restore(path, RENAMES_H5_DATASET, record) + + runs = [] + with ( + patch.object(legacy_inputs, "LEGACY_INPUT_RENAMES", register), + patch.object(pe.Simulation, "run", lambda self: runs.append(self.id)), + patch.object(pe.Simulation, "save", lambda self: None), + ): + problem = output_stamp_problem(path) + loading = simulation() + if problem is None: + loading.load() + assert _person_outputs(loading)["wic"].iloc[INFANT] == 0 + else: + with pytest.raises(ValueError) as refused: + loading.load() + assert str(refused.value) == OUT_OF_DATE.format( + path=path, reason=problem + ) + + _cache.clear() + try: + simulation().ensure() + finally: + _cache.clear() + + assert bool(runs) == (problem is not None) + event("current" if problem is None else "out of date") diff --git a/tests/test_us_output_stamp.py b/tests/test_us_output_stamp.py new file mode 100644 index 00000000..25ae2b37 --- /dev/null +++ b/tests/test_us_output_stamp.py @@ -0,0 +1,1079 @@ +"""Unit and property tests for the stamp on saved US simulation outputs. + +A US run records a stamp in its output (``legacy_inputs.OUTPUT_STAMP_H5_DATASET`` +once saved): the register of renamed stored inputs it was calculated under, an +output format revision, and the stamp of the year file its input data was cut +into, or ``null``. ``Simulation.load()`` refuses an output, and +``Simulation.ensure()`` calculates it again, unless the stamp matches this +release's. These tests state the check's invariants over generated stamps, +records and registers, without a country engine: + +- **Exactness.** An output is current exactly when its stamp has this + release's output format revision and register, its year-file stamp is + ``null`` or this release's year-file stamp, and its renames record lists + only renames in its stamped register. +- **Totality.** Whatever the stamp and record hold, the check returns a + reason or ``None``, and raises only for errors unrelated to the file's + contents (a missing file, no permission, a directory, a lock). +- **Round trip.** A stamp reads back as written, and writing again replaces it. +- **Named renames.** The reason names every rename the output's registers + (its own and its year file's) and this release's disagree on. +- **No partial output is current.** Every truncation of a current output is + out of date. +- **Distinct stamps.** An output never passes as a year file, nor a year file + as an output. + +``test_us_legacy_inputs_integration.py`` saves, loads and ensures outputs over +a real policyengine-us engine. +""" + +from __future__ import annotations + +import functools +import json +import os +import tempfile +from pathlib import Path +from unittest.mock import patch + +import h5py +import numpy as np +import pandas as pd +import pytest +from hypothesis import event, given, settings +from hypothesis import strategies as st +from microdf import MicroDataFrame + +from policyengine.tax_benefit_models.us import legacy_inputs +from policyengine.tax_benefit_models.us.datasets import ( + PolicyEngineUSDataset, + USYearData, +) +from policyengine.tax_benefit_models.us.legacy_inputs import ( + LEGACY_INPUT_RENAMES, + OUTPUT_STAMP_H5_DATASET, + OUTPUT_STAMP_REVISION, + RENAMES_H5_DATASET, + YEAR_FILE_STAMP_H5_DATASET, + YEAR_FILE_STAMP_KEY, + YEAR_FILE_STAMP_REVISION, + output_stamp, + output_stamp_problem, + read_output_stamp, + read_renames_record, + read_year_file_stamp, + write_output_stamp, + year_file_stamp, + year_file_stamp_metadata, + year_file_stamp_problem, +) + +LEGACY = "would_claim_wic" +LIVE = "takes_up_wic_if_eligible" +REVISION = OUTPUT_STAMP_REVISION +YEAR_FILE_REVISION = YEAR_FILE_STAMP_REVISION +GROUPS = ("household", "tax_unit", "spm_unit", "family", "marital_unit") + +#: A small alphabet, so generated registers often share names and disagree +#: only on some entries. +NAMES = st.sampled_from(["alpha", "beta", "gamma", "delta"]) +REGISTERS = st.dictionaries(NAMES, NAMES, max_size=3) +#: Raw node contents: ``_ABSENT`` for no node, a JSON value, or a node the +#: reader cannot parse as JSON. +_ABSENT = object() +JSON_VALUES = st.recursive( + st.none() | st.booleans() | st.integers() | st.floats() | st.text(max_size=8), + lambda children: ( + st.lists(children, max_size=3) + | st.dictionaries(st.text(max_size=8), children, max_size=3) + ), + max_leaves=6, +) +#: Text stored as it is, not as JSON. HDF5 variable-length strings cannot +#: hold a NUL character, and a lone surrogate cannot be encoded as UTF-8. +RAW_TEXT = st.text( + alphabet=st.characters(blacklist_categories=("Cs",), blacklist_characters="\x00"), + max_size=12, +) + + +def _frames() -> dict[str, pd.DataFrame]: + person = pd.DataFrame( + { + "person_id": [1, 2], + **{f"{group}_id": [1, 1] for group in GROUPS}, + "person_weight": [1.0, 1.0], + "wic": [0.0, 50.0], + } + ) + frames = {"person": person} + for group in GROUPS: + frames[group] = pd.DataFrame({f"{group}_id": [1], f"{group}_weight": [1.0]}) + return frames + + +def _dataset(path) -> PolicyEngineUSDataset: + return PolicyEngineUSDataset( + name="output", + description="Uncertified two-person output fixture", + filepath=str(path), + year=2024, + data=USYearData( + **{ + entity: MicroDataFrame(frame, weights=f"{entity}_weight") + for entity, frame in _frames().items() + } + ), + ) + + +def _write_node(file, name, content) -> None: + """Store ``content`` under ``name``: JSON, or a node kind to test.""" + if content is _ABSENT: + return + if isinstance(content, tuple): + kind, value = content + if kind == "text": + file.create_dataset(name, data=value, dtype=h5py.string_dtype("utf-8")) + elif kind == "bytes": + file.create_dataset(name, data=np.bytes_(value)) + elif kind == "number": + file.create_dataset(name, data=value) + elif kind == "array": + file.create_dataset(name, data=value, dtype=h5py.string_dtype("utf-8")) + elif kind == "group": + file.create_group(name) + else: + raise AssertionError(kind) + return + file.create_dataset( + name, data=json.dumps(content), dtype=h5py.string_dtype("utf-8") + ) + + +def _write_output(path, *, stamp=_ABSENT, record=_ABSENT, year_file_stamp=_ABSENT): + """Write an output's tables, then its record and stamps as given. + + ``stamp`` (the output stamp), ``record`` and ``year_file_stamp`` are JSON + values, ``_ABSENT``, or a raw node (see :func:`_write_node`). + """ + Path(path).parent.mkdir(parents=True, exist_ok=True) + _dataset(path).save() + with h5py.File(path, "a") as file: + _write_node(file, RENAMES_H5_DATASET, record) + _write_node(file, YEAR_FILE_STAMP_H5_DATASET, year_file_stamp) + _write_node(file, OUTPUT_STAMP_H5_DATASET, stamp) + return Path(path) + + +def _same_json(left, right) -> bool: + """Equal as JSON text, so ``true`` is not ``1`` and ``1.0`` is not ``1``.""" + return json.dumps(left, sort_keys=True) == json.dumps(right, sort_keys=True) + + +def _str_map(value) -> bool: + return isinstance(value, dict) and all( + isinstance(key, str) and isinstance(item, str) for key, item in value.items() + ) + + +def _is_current(stamp, record, register) -> bool: + """The exactness invariant, stated directly on the stored JSON values.""" + year_file = {"register": dict(register), "revision": YEAR_FILE_REVISION} + return ( + isinstance(stamp, dict) + and set(stamp) == {"register", "revision", "year_file"} + and type(stamp["revision"]) is int + and stamp["revision"] == REVISION + and _str_map(stamp["register"]) + and stamp["register"] == register + and (stamp["year_file"] is None or _same_json(stamp["year_file"], year_file)) + and _str_map(record) + and record.items() <= stamp["register"].items() + ) + + +def _stamp(register, revision=REVISION, year_file=None): + return {"register": dict(register), "revision": revision, "year_file": year_file} + + +def _year_stamp(register, revision=YEAR_FILE_REVISION): + return {"register": dict(register), "revision": revision} + + +# --- The stamp ---------------------------------------------------------- + + +def test_the_stamp_holds_the_register_the_revision_and_the_year_file(): + assert output_stamp() == { + "register": {LEGACY: LIVE}, + "revision": 1, + "year_file": None, + } + assert output_stamp()["register"] == LEGACY_INPUT_RENAMES + assert output_stamp(year_file_stamp()) == { + "register": {LEGACY: LIVE}, + "revision": 1, + "year_file": {"register": {LEGACY: LIVE}, "revision": 1}, + } + + +def test_the_stamp_reads_the_register_and_revision_when_called(monkeypatch): + """model.py must not freeze the register at import.""" + monkeypatch.setattr(legacy_inputs, "LEGACY_INPUT_RENAMES", {"b": "c", "a": "d"}) + monkeypatch.setattr(legacy_inputs, "OUTPUT_STAMP_REVISION", 4) + + assert output_stamp() == { + "register": {"a": "d", "b": "c"}, + "revision": 4, + "year_file": None, + } + + +def test_the_stamp_copies_the_year_file_stamp(): + """A later change to the input's metadata does not change the stamp.""" + year_file = year_file_stamp() + stamp = output_stamp(year_file) + year_file["register"]["x"] = "y" + year_file["revision"] = 9 + + assert stamp["year_file"] == {"register": {LEGACY: LIVE}, "revision": 1} + + +def test_the_stamp_is_stored_as_utf8_json(tmp_path): + path = _write_output(tmp_path / "output.h5") + + write_output_stamp(path, output_stamp()) + + with h5py.File(path, "r") as file: + assert file[OUTPUT_STAMP_H5_DATASET].asstr()[()] == ( + f'{{"register": {{"{LEGACY}": "{LIVE}"}}, "revision": 1, ' + '"year_file": null}' + ) + + +@settings(max_examples=50, deadline=None) +@given( + first=st.fixed_dictionaries( + {"register": REGISTERS, "revision": st.integers(), "year_file": JSON_VALUES} + ), + second=st.fixed_dictionaries( + {"register": REGISTERS, "revision": st.integers(), "year_file": JSON_VALUES} + ), +) +def test_an_output_stamp_reads_back_as_written(first, second): + """Round trip: a stamp reads back exactly, a new one replaces it, and + the tables, the record and any year-file stamp are left alone.""" + with tempfile.TemporaryDirectory() as directory: + path = _write_output( + Path(directory) / "output.h5", record={}, year_file_stamp=[1] + ) + assert read_output_stamp(path) is None + + write_output_stamp(path, first) + assert _same_json(read_output_stamp(path), first) + + write_output_stamp(path, second) + assert _same_json(read_output_stamp(path), second) + assert read_renames_record(path) == {} + assert read_year_file_stamp(path) == [1] + with pd.HDFStore(path, mode="r") as store: + assert store["person"]["wic"].tolist() == [0.0, 50.0] + + +def test_a_missing_or_foreign_file_has_no_output_stamp(tmp_path): + assert read_output_stamp(tmp_path / "missing.h5") is None + text = tmp_path / "notes.txt" + text.write_text("not an H5 file") + assert read_output_stamp(text) is None + + +# --- Exactness ---------------------------------------------------------- + + +@st.composite +def near(draw, register): + """``register`` itself (most often), one entry away from it, or any.""" + exact = st.just(dict(register)) + return draw( + st.one_of( + exact, + exact, + exact, + REGISTERS, + st.builds( + lambda name: {k: v for k, v in register.items() if k != name}, + NAMES, + ), + st.builds(lambda k, v: {**register, k: v}, NAMES, NAMES), + ) + ) + + +@st.composite +def stamped_outputs(draw): + """A current register, and an output's stamp and record near it. + + Each part is usually right, so current outputs and outputs wrong in a + single part are both common. + """ + register = draw(REGISTERS) + stamped = draw(near(register)) + revision = draw(st.sampled_from([REVISION] * 6 + [0, 2, -1, True, 1.0])) + year_file = draw( + st.one_of( + st.none(), + st.none(), + st.builds( + _year_stamp, + near(register), + st.sampled_from([YEAR_FILE_REVISION] * 6 + [0, 2, True, 1.0, "1"]), + ), + st.builds(_year_stamp, near(register)), + st.just("unreadable"), + JSON_VALUES, + ) + ) + subset = st.builds( + lambda keep: {k: v for k, v in stamped.items() if k in keep}, + st.sets(NAMES), + ) + record = draw( + # A subset of the stamped register, as a run records, most often. + st.one_of(subset, subset, subset, st.just(_ABSENT), REGISTERS, JSON_VALUES) + ) + return register, _stamp(stamped, revision, year_file), record + + +@settings(max_examples=300, deadline=None) +@given(case=stamped_outputs()) +def test_an_output_is_current_only_under_this_releases_exact_stamp(case): + """Exactness, and refinement of the record check outputs had before. + + Equivalently: the stamp is what this release would write for a run over + a dataset with no year-file stamp, or over a current year file. + """ + register, stamp, record = case + with tempfile.TemporaryDirectory() as directory: + path = _write_output(Path(directory) / "output.h5", stamp=stamp, record=record) + with patch.object(legacy_inputs, "LEGACY_INPUT_RENAMES", register): + problem = output_stamp_problem(path) + expected = ( + any( + _same_json(stamp, output_stamp(year_file)) + for year_file in (None, year_file_stamp()) + ) + and _str_map(record) + and record.items() <= stamp["register"].items() + ) + + assert (problem is None) == _is_current(stamp, record, register) == expected + assert problem is None or (isinstance(problem, str) and problem) + event( + "current" + if problem is None + else "out of date: " + ("year file" if "year file" in problem else "other") + ) + if problem is None: + # Every output the stamp accepts, the record check accepted too. + assert read_renames_record(path) == record + assert _same_json(read_output_stamp(path), stamp) + + +NODES = st.one_of( + st.just(_ABSENT), + JSON_VALUES, + st.fixed_dictionaries( + { + "register": st.one_of(REGISTERS, JSON_VALUES), + "revision": st.one_of( + st.just(REVISION), st.booleans(), st.floats(), JSON_VALUES + ), + "year_file": st.one_of( + st.none(), + st.fixed_dictionaries( + { + "register": st.one_of(REGISTERS, JSON_VALUES), + "revision": st.one_of( + st.just(YEAR_FILE_REVISION), st.booleans(), JSON_VALUES + ), + } + ), + JSON_VALUES, + ), + } + ), + st.tuples(st.just("text"), RAW_TEXT), + st.tuples(st.just("bytes"), st.binary(min_size=1, max_size=12)), + st.tuples(st.just("number"), st.integers(-5, 5)), + st.tuples(st.just("array"), st.lists(RAW_TEXT, min_size=1, max_size=3)), + st.tuples(st.just("group"), st.none()), +) + + +@settings(max_examples=300, deadline=None) +@given(stamp=NODES, record=NODES, register=REGISTERS) +def test_the_output_problem_is_total_and_exact_for_any_content(stamp, record, register): + """Totality: any stamp and record give a reason or ``None``, never an + error, and ``None`` only for this release's exact stamp.""" + with tempfile.TemporaryDirectory() as directory: + path = _write_output(Path(directory) / "output.h5", stamp=stamp, record=record) + with patch.object(legacy_inputs, "LEGACY_INPUT_RENAMES", register): + problem = output_stamp_problem(path) + + assert problem is None or (isinstance(problem, str) and problem) + if problem is None: + assert _is_current(stamp, record, register) + + +@pytest.mark.parametrize( + "stamp", + [ + {"register": {LEGACY: LIVE}, "revision": True, "year_file": None}, + {"register": {LEGACY: LIVE}, "revision": 1.0, "year_file": None}, + {"register": {LEGACY: LIVE}, "revision": "1", "year_file": None}, + {"register": {LEGACY: LIVE}, "revision": 1}, + {"register": {LEGACY: LIVE}, "year_file": None}, + {"register": {LEGACY: LIVE}, "revision": 1, "year_file": None, "x": None}, + {"register": {LEGACY: 5}, "revision": 1, "year_file": None}, + {"register": [[LEGACY, LIVE]], "revision": 1, "year_file": None}, + [{"register": {LEGACY: LIVE}, "revision": 1, "year_file": None}], + None, + ("text", "not json"), + ("text", ""), + ("bytes", b"\xff\xfe"), + ("number", 1), + ("array", ["{}"]), + ("group", None), + ], +) +def test_a_malformed_stamp_is_out_of_date(tmp_path, stamp): + path = _write_output(tmp_path / "output.h5", stamp=stamp, record={}) + + assert output_stamp_problem(path) == ( + "its stamp of the register of renamed inputs it was calculated under is " + "malformed" + ) + + +@pytest.mark.parametrize("revision", [0, 2, -1]) +def test_an_output_from_another_format_revision_is_out_of_date(tmp_path, revision): + # A later revision may shape the rest of its stamp and record differently, + # and is still reported by revision rather than as malformed. + stamp = {"revision": revision, "registers": [[LEGACY, LIVE]]} + path = _write_output(tmp_path / "output.h5", stamp=stamp) + + assert output_stamp_problem(path) == ( + f"it was saved with output format revision {revision}, and this " + "version of policyengine.py reads only revision 1" + ) + + +def test_outputs_without_a_stamp_are_out_of_date_with_distinct_reasons(tmp_path): + """Outputs saved before the mapping wrote neither; #531 wrote the record.""" + oldest = _write_output(tmp_path / "oldest.h5") + record_only = _write_output(tmp_path / "record_only.h5", record={LEGACY: LIVE}) + + assert output_stamp_problem(oldest) == ( + "it has no record of the renamed stored inputs mapped when it was " + "calculated, so it was saved by an older policyengine.py or not saved " + f"completely, and may have been calculated without any of: {LEGACY} -> " + f"{LIVE}" + ) + assert output_stamp_problem(record_only) == ( + "it has no stamp of the register of renamed inputs it was calculated " + f"under, so it may have been calculated without any of: {LEGACY} -> {LIVE}" + ) + + +def test_an_unstamped_output_is_out_of_date_even_with_an_empty_register( + tmp_path, monkeypatch +): + monkeypatch.setattr(legacy_inputs, "LEGACY_INPUT_RENAMES", {}) + oldest = _write_output(tmp_path / "oldest.h5") + record_only = _write_output(tmp_path / "record_only.h5", record={}) + + assert output_stamp_problem(oldest) == ( + "it has no record of the renamed stored inputs mapped when it was " + "calculated, so it was saved by an older policyengine.py or not saved " + "completely" + ) + assert output_stamp_problem(record_only) == ( + "it has no stamp of the register of renamed inputs it was calculated under" + ) + + +MALFORMED_RECORD = ( + "its record of the renamed stored inputs mapped when it was calculated is malformed" +) + + +@pytest.mark.parametrize( + "record, reason", + [ + ( + _ABSENT, + "it has no record of the renamed stored inputs mapped when it was " + "calculated", + ), + ([], MALFORMED_RECORD), + ([[LEGACY, LIVE]], MALFORMED_RECORD), + (5, MALFORMED_RECORD), + (None, MALFORMED_RECORD), + ({LEGACY: 1}, MALFORMED_RECORD), + (("text", "{"), MALFORMED_RECORD), + (("group", None), MALFORMED_RECORD), + ( + {"retired": "replacement"}, + "its record lists renames outside the register it was stamped with: " + "retired -> replacement", + ), + ( + {LEGACY: "elsewhere"}, + "its record lists renames outside the register it was stamped with: " + f"{LEGACY} -> elsewhere", + ), + ], +) +def test_a_stamped_output_needs_a_record_within_its_register(tmp_path, record, reason): + path = _write_output(tmp_path / "output.h5", stamp=output_stamp(), record=record) + + assert output_stamp_problem(path) == reason + + +@pytest.mark.parametrize("record", [{}, {LEGACY: LIVE}]) +@pytest.mark.parametrize("year_file", [None, year_file_stamp()]) +def test_an_output_stamped_with_this_releases_stamp_is_current( + tmp_path, record, year_file +): + path = _write_output( + tmp_path / "output.h5", stamp=output_stamp(year_file), record=record + ) + + assert output_stamp_problem(path) is None + + +def test_a_stamp_and_record_stored_as_fixed_length_strings_are_read(tmp_path): + stamp = json.dumps(output_stamp(year_file_stamp())).encode() + record = json.dumps({LEGACY: LIVE}).encode() + path = _write_output( + tmp_path / "output.h5", stamp=("bytes", stamp), record=("bytes", record) + ) + + assert output_stamp_problem(path) is None + assert read_output_stamp(path) == output_stamp(year_file_stamp()) + + +# --- The year file an output was calculated from ------------------------ + + +MALFORMED_YEAR_FILE = ( + "its record of the stamp of the year file it was calculated from is malformed" +) + + +@pytest.mark.parametrize( + "year_file, reason", + [ + ( + _year_stamp({}), + "it was calculated from a year file cut before these renames were " + f"registered, which may have lost them: {LEGACY} -> {LIVE}", + ), + ( + _year_stamp({**LEGACY_INPUT_RENAMES, "old": "new"}), + "it was calculated from a year file cut under renames this version of " + "policyengine.py does not register: old -> new", + ), + ( + _year_stamp({LEGACY: "elsewhere"}), + "it was calculated from a year file cut before these renames were " + f"registered, which may have lost them: {LEGACY} -> {LIVE}; it was " + "calculated from a year file cut under renames this version of " + f"policyengine.py does not register: {LEGACY} -> elsewhere", + ), + ( + _year_stamp(LEGACY_INPUT_RENAMES, 0), + "it was calculated from a year file cut with year-file format " + "revision 0, and this version of policyengine.py reads only revision 1", + ), + ( + {"revision": 2, "registers": []}, + "it was calculated from a year file cut with year-file format " + "revision 2, and this version of policyengine.py reads only revision 1", + ), + (_year_stamp(LEGACY_INPUT_RENAMES, True), MALFORMED_YEAR_FILE), + (_year_stamp(LEGACY_INPUT_RENAMES, 1.0), MALFORMED_YEAR_FILE), + ({**year_file_stamp(), "extra": None}, MALFORMED_YEAR_FILE), + ({"register": {LEGACY: LIVE}}, MALFORMED_YEAR_FILE), + ({"register": {LEGACY: 1}, "revision": 1}, MALFORMED_YEAR_FILE), + ("unreadable", MALFORMED_YEAR_FILE), + (False, MALFORMED_YEAR_FILE), + (0, MALFORMED_YEAR_FILE), + ([], MALFORMED_YEAR_FILE), + ({}, MALFORMED_YEAR_FILE), + ], +) +def test_an_output_from_an_out_of_date_year_file_is_out_of_date( + tmp_path, year_file, reason +): + """A run over a year file maps only what the file still stores under a + legacy name, so the output inherits whatever the file lost.""" + path = _write_output( + tmp_path / "output.h5", + stamp=_stamp(LEGACY_INPUT_RENAMES, 1, year_file), + record={}, + ) + + assert output_stamp_problem(path) == reason + + +def test_raising_the_year_file_revision_makes_outputs_from_year_files_out_of_date( + tmp_path, monkeypatch +): + from_year_file = _write_output( + tmp_path / "from_year_file.h5", stamp=output_stamp(year_file_stamp()), record={} + ) + from_source = _write_output( + tmp_path / "from_source.h5", stamp=output_stamp(), record={} + ) + monkeypatch.setattr(legacy_inputs, "YEAR_FILE_STAMP_REVISION", 2) + + assert output_stamp_problem(from_year_file) == ( + "it was calculated from a year file cut with year-file format revision 1, " + "and this version of policyengine.py reads only revision 2" + ) + assert output_stamp_problem(from_source) is None + + +def test_raising_the_output_revision_makes_older_outputs_out_of_date( + tmp_path, monkeypatch +): + older = _write_output(tmp_path / "older.h5", stamp=output_stamp(), record={}) + monkeypatch.setattr(legacy_inputs, "OUTPUT_STAMP_REVISION", 7) + newer = _write_output(tmp_path / "newer.h5", stamp=output_stamp(), record={}) + + assert output_stamp()["revision"] == 7 + assert output_stamp_problem(newer) is None + assert output_stamp_problem(older) == ( + "it was saved with output format revision 1, and this version of " + "policyengine.py reads only revision 7" + ) + + +# --- Named renames ------------------------------------------------------ + + +def _expected_differences(stamped, register): + missing = {k: v for k, v in register.items() if stamped.get(k) != v} + unknown = {k: v for k, v in stamped.items() if register.get(k) != v} + return missing, unknown + + +@settings(max_examples=200, deadline=None) +@given(register=REGISTERS, stamped=REGISTERS, year_file=REGISTERS) +def test_a_stale_output_names_each_rename_it_may_lack(register, stamped, year_file): + """The reason names every entry its registers and this release's disagree + on, as ``old -> new`` from the register that has it, and no other.""" + with tempfile.TemporaryDirectory() as directory: + path = _write_output( + Path(directory) / "output.h5", + stamp=_stamp(stamped, year_file=_year_stamp(year_file)), + record={}, + ) + with patch.object(legacy_inputs, "LEGACY_INPUT_RENAMES", register): + problem = output_stamp_problem(path) + + if stamped == register and year_file == register: + assert problem is None + return + reasons = problem.split("; ") + own = [r for r in reasons if "from a year file" not in r] + cut = [r for r in reasons if "from a year file" in r] + for registered, clauses in ((stamped, own), (year_file, cut)): + missing, unknown = _expected_differences(registered, register) + assert len(clauses) == bool(missing) + bool(unknown) + if missing: + assert "before these renames were registered" in clauses[0] + for legacy, live in missing.items(): + assert f"{legacy} -> {live}" in clauses[0] + if unknown: + assert "does not register" in clauses[-1] + for legacy, live in unknown.items(): + assert f"{legacy} -> {live}" in clauses[-1] + for legacy, live in register.items() & registered.items(): + assert all(f"{legacy} -> {live}" not in clause for clause in clauses) + + +def test_a_future_register_entry_makes_a_stamped_output_out_of_date( + tmp_path, monkeypatch +): + path = _write_output(tmp_path / "output.h5", stamp=output_stamp(), record={}) + monkeypatch.setattr( + legacy_inputs, + "LEGACY_INPUT_RENAMES", + {**LEGACY_INPUT_RENAMES, "old_input": "new_input"}, + ) + + assert output_stamp_problem(path) == ( + "it was calculated before these renames were registered, so it may have " + "been calculated without them: old_input -> new_input" + ) + + +@pytest.mark.parametrize( + "register, reason", + [ + ( + {}, + "it was calculated under renames this version of policyengine.py does " + f"not register: {LEGACY} -> {LIVE}", + ), + ( + {LEGACY: "elsewhere"}, + "it was calculated before these renames were registered, so it may " + f"have been calculated without them: {LEGACY} -> elsewhere; it was " + "calculated under renames this version of policyengine.py does not " + f"register: {LEGACY} -> {LIVE}", + ), + ], +) +def test_retiring_or_retargeting_an_entry_makes_a_stamped_output_out_of_date( + tmp_path, monkeypatch, register, reason +): + path = _write_output( + tmp_path / "output.h5", stamp=output_stamp(), record={LEGACY: LIVE} + ) + monkeypatch.setattr(legacy_inputs, "LEGACY_INPUT_RENAMES", register) + + assert output_stamp_problem(path) == reason + + +# --- No partial output is current --------------------------------------- + + +@functools.cache +def _current_output_bytes() -> bytes: + with tempfile.TemporaryDirectory() as directory: + path = _write_output( + Path(directory) / "output.h5", + stamp=output_stamp(year_file_stamp()), + record={LEGACY: LIVE}, + ) + assert output_stamp_problem(path) is None + return path.read_bytes() + + +@settings(max_examples=100, deadline=None) +@given(data=st.data()) +def test_a_truncated_output_is_never_current(data): + content = _current_output_bytes() + length = data.draw(st.integers(0, len(content) - 1)) + with tempfile.TemporaryDirectory() as directory: + path = Path(directory) / "output.h5" + path.write_bytes(content[:length]) + + assert output_stamp_problem(path) is not None + + +def test_a_file_that_is_not_hdf5_is_out_of_date(tmp_path): + path = tmp_path / "output.h5" + path.write_text("not an H5 file") + + assert output_stamp_problem(path) == ( + "it cannot be read as an HDF5 file (it may be half-written or corrupt)" + ) + + +def test_saving_the_tables_again_drops_the_stamp(tmp_path): + """``PolicyEngineUSDataset.save()`` truncates; only the model version's + ``save()`` stamps an output.""" + path = _write_output(tmp_path / "output.h5", stamp=output_stamp(), record={}) + + _dataset(path).save() + + assert read_output_stamp(path) is None + assert output_stamp_problem(path) is not None + + +#: JSON nested very deeply. Whether the parser reaches its limit depends on +#: the interpreter and the platform (see ``test_us_year_file_stamp.py``). +DEEP_JSON = "[" * 100_000 + "0" + "]" * 100_000 + + +@pytest.mark.parametrize( + "stamp, record, reason", + [ + ( + ("text", DEEP_JSON), + {}, + "its stamp of the register of renamed inputs it was calculated under " + "is malformed", + ), + (output_stamp(), ("text", DEEP_JSON), MALFORMED_RECORD), + ], +) +def test_json_nested_too_deeply_is_out_of_date_not_an_error( + tmp_path, stamp, record, reason +): + path = _write_output(tmp_path / "output.h5", stamp=stamp, record=record) + + assert output_stamp_problem(path) == reason + + +@pytest.mark.parametrize("name", [OUTPUT_STAMP_H5_DATASET, RENAMES_H5_DATASET]) +def test_a_recursion_error_while_parsing_is_unreadable_content( + tmp_path, monkeypatch, name +): + """The RecursionError path, on every platform.""" + path = _write_output(tmp_path / "output.h5", stamp=output_stamp(), record={}) + loads = json.loads + + def deep(text, *args, **kwargs): + with h5py.File(path, "r") as file: + if text == file[name].asstr()[()]: + raise RecursionError("maximum recursion depth exceeded") + return loads(text, *args, **kwargs) + + monkeypatch.setattr(legacy_inputs.json, "loads", deep) + + assert output_stamp_problem(path) == ( + "its stamp of the register of renamed inputs it was calculated under is " + "malformed" + if name == OUTPUT_STAMP_H5_DATASET + else MALFORMED_RECORD + ) + + +def test_dangling_links_in_place_of_the_stamp_and_record_read_as_absent(tmp_path): + path = _write_output(tmp_path / "output.h5") + with h5py.File(path, "a") as file: + file[OUTPUT_STAMP_H5_DATASET] = h5py.SoftLink("/nowhere") + file[RENAMES_H5_DATASET] = h5py.ExternalLink( + str(tmp_path / "missing.h5"), "/record" + ) + + assert output_stamp_problem(path).startswith( + "it has no record of the renamed stored inputs mapped when it was calculated" + ) + assert read_output_stamp(path) is None + + +# --- Errors unrelated to the file's contents ---------------------------- + + +def test_a_missing_output_raises_file_not_found(tmp_path): + """``Simulation.ensure()`` treats it as a plain miss.""" + with pytest.raises(FileNotFoundError): + output_stamp_problem(tmp_path / "missing.h5") + + +def _unreadable_by_permission(path) -> Path: + path = _write_output(path, stamp=output_stamp(), record={}) + os.chmod(path, 0) + return path + + +def _a_directory(path) -> Path: + path.mkdir() + return path + + +@pytest.mark.parametrize( + "make, error", + [ + (_unreadable_by_permission, PermissionError), + (_a_directory, IsADirectoryError), + ], +) +def test_errors_unrelated_to_the_contents_are_raised(tmp_path, make, error): + if error is PermissionError and hasattr(os, "geteuid") and os.geteuid() == 0: + pytest.skip("root can read a file without permission bits") + path = make(tmp_path / "output.h5") + try: + with pytest.raises(error): + output_stamp_problem(path) + with pytest.raises(error): + year_file_stamp_metadata(path) + finally: + if path.is_file(): + os.chmod(path, 0o644) + + +@pytest.mark.parametrize("check", [output_stamp_problem, year_file_stamp_metadata]) +def test_a_locked_file_is_raised(tmp_path, monkeypatch, check): + path = _write_output(tmp_path / "output.h5", stamp=output_stamp(), record={}) + + def locked(*args, **kwargs): + raise BlockingIOError(35, "Unable to lock file") + + monkeypatch.setattr(h5py, "File", locked) + with pytest.raises(BlockingIOError): + check(path) + + +@pytest.mark.parametrize( + "error", + [ + PermissionError(13, "Permission denied"), + IsADirectoryError(21, "Is a directory"), + BlockingIOError(35, "Unable to lock file"), + ], +) +def test_errors_unrelated_to_the_contents_are_raised_while_reading( + tmp_path, monkeypatch, error +): + path = _write_output( + tmp_path / "output.h5", + stamp=output_stamp(), + record={}, + year_file_stamp=year_file_stamp(), + ) + + def fail(self, *args, **kwargs): + raise error + + monkeypatch.setattr(h5py.Dataset, "asstr", fail) + with pytest.raises(type(error)): + output_stamp_problem(path) + with pytest.raises(type(error)): + year_file_stamp_metadata(path) + + +def test_the_decision_reads_only_the_stamp_and_the_record(tmp_path, monkeypatch): + path = _write_output( + tmp_path / "output.h5", stamp=output_stamp(year_file_stamp()), record={} + ) + real_file = h5py.File + + class StampAndRecordOnly: + def __init__(self, path, mode="r", **kwargs): + assert mode == "r" + self._file = real_file(path, mode, **kwargs) + + def __enter__(self): + return self + + def __exit__(self, *exc_info): + self._file.close() + + def get(self, name, default=None): + assert name in (OUTPUT_STAMP_H5_DATASET, RENAMES_H5_DATASET), name + return self._file.get(name, default) + + monkeypatch.setattr(h5py, "File", StampAndRecordOnly) + + assert output_stamp_problem(path) is None + + +# --- Distinct stamps: an output never passes as a year file ------------- + + +@settings(max_examples=100, deadline=None) +@given( + register=REGISTERS, + keep=st.sets(NAMES), + from_year_file=st.booleans(), +) +def test_an_output_never_passes_as_a_year_file_nor_the_reverse( + register, keep, from_year_file +): + """Current files of each kind, under any register, are out of date as the + other kind, and so is either stamp stored under the other's name.""" + record = {k: v for k, v in register.items() if k in keep} + with ( + tempfile.TemporaryDirectory() as directory, + patch.object(legacy_inputs, "LEGACY_INPUT_RENAMES", register), + ): + stamp = output_stamp(year_file_stamp() if from_year_file else None) + output = _write_output( + Path(directory) / "output.h5", stamp=stamp, record=record + ) + year_file = _write_output( + Path(directory) / "year.h5", + year_file_stamp=year_file_stamp(), + record=record, + ) + assert output_stamp_problem(output) is None + assert year_file_stamp_problem(year_file) is None + + assert year_file_stamp_problem(output) is not None + assert output_stamp_problem(year_file) is not None + + # Each stamp stored under the other's name. + swapped_output = _write_output( + Path(directory) / "swapped_output.h5", + year_file_stamp=stamp, + record=record, + ) + swapped_year_file = _write_output( + Path(directory) / "swapped_year.h5", + stamp=year_file_stamp(), + record=record, + ) + assert year_file_stamp_problem(swapped_output) is not None + assert output_stamp_problem(swapped_year_file) is not None + + +def test_the_two_stamps_have_different_keys(): + assert set(output_stamp()) == {"register", "revision", "year_file"} + assert set(year_file_stamp()) == {"register", "revision"} + assert OUTPUT_STAMP_H5_DATASET != YEAR_FILE_STAMP_H5_DATASET + + +# --- The year-file stamp a dataset keeps -------------------------------- + + +@pytest.mark.parametrize( + "stored, kept", + [ + (_ABSENT, None), + ({"register": {LEGACY: LIVE}, "revision": 1}, year_file_stamp()), + ([1, 2], [1, 2]), + (False, False), + (0, 0), + ("text", "text"), + (None, "unreadable"), + (("text", "not json"), "unreadable"), + (("bytes", b"\xff"), "unreadable"), + (("number", 3), "unreadable"), + (("group", None), "unreadable"), + ], +) +def test_a_dataset_keeps_its_year_files_stamp(tmp_path, stored, kept): + path = _write_output(tmp_path / "year.h5", year_file_stamp=stored, record={}) + + assert year_file_stamp_metadata(path) == kept + dataset = PolicyEngineUSDataset( + name="year-file", description="reload", filepath=str(path), year=2024 + ) + if kept is None: + assert YEAR_FILE_STAMP_KEY not in dataset.metadata + else: + assert dataset.metadata[YEAR_FILE_STAMP_KEY] == kept + + +def test_a_file_that_is_not_hdf5_keeps_no_year_file_stamp(tmp_path): + path = tmp_path / "notes.h5" + path.write_text("not an H5 file") + + assert year_file_stamp_metadata(path) is None + + +def test_reloading_a_resaved_year_file_drops_the_stamp_it_kept(tmp_path): + path = _write_output( + tmp_path / "year.h5", year_file_stamp=year_file_stamp(), record={} + ) + dataset = PolicyEngineUSDataset( + name="year-file", description="reload", filepath=str(path), year=2024 + ) + assert dataset.metadata[YEAR_FILE_STAMP_KEY] == year_file_stamp() + + # Saving drops the file's stamp, so the data reloaded from it is no + # longer a stamped year file's. + dataset.save() + dataset.load() + + assert read_year_file_stamp(path) is None + assert YEAR_FILE_STAMP_KEY not in dataset.metadata + assert dataset.metadata["legacy_input_renames"] == {} From e726177973a579a42df093f28262d731f66fc8b4 Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Tue, 13 Oct 2026 07:31:26 -0400 Subject: [PATCH 3/6] Address review 1: cyclic links, oracle decoding, scoped provenance - The shared H5 JSON reader treats RuntimeError, which h5py raises for a soft link that leads back to itself and which covers RecursionError, as unreadable content, so output_stamp_problem, year_file_stamp_problem and year_file_stamp_metadata return a reason rather than raise. - The totality property decodes raw nodes independently and asserts both directions; the exactness generator weights with sampled_from, since st.one_of drops repeated branches. - A region-scoped run records the year-file stamp of the dataset it was given, tested over a current and an out-of-date year file. - The changelog no longer overstates the stale-year-file guarantee. Co-Authored-By: Claude Opus 5.5 --- changelog.d/592.fixed.md | 2 +- .../tax_benefit_models/us/legacy_inputs.py | 4 +- tests/test_us_legacy_inputs_integration.py | 41 +++++ tests/test_us_output_stamp.py | 140 ++++++++++++++---- 4 files changed, 158 insertions(+), 29 deletions(-) diff --git a/changelog.d/592.fixed.md b/changelog.d/592.fixed.md index 204b811e..786a5f47 100644 --- a/changelog.d/592.fixed.md +++ b/changelog.d/592.fixed.md @@ -1 +1 @@ -A US `Simulation.run()` stamps its output with the register of renamed stored inputs it was calculated under, an output format revision, and the stamp of the year file its input data was cut into, if any. `Simulation.load()` refuses a saved US output, and `Simulation.ensure()` runs it again, unless that stamp matches this release exactly. An output calculated before a rename was registered, or from a year file that is out of date, is therefore not reused. The error says why, naming the renames the output may lack. Outputs saved before the stamp existed are calculated again once. +A US `Simulation.run()` stamps its output with the register of renamed stored inputs it was calculated under, an output format revision, and the stamp of the year file its input data was cut into, if any. `Simulation.load()` refuses a saved US output, and `Simulation.ensure()` runs it again, unless that stamp matches this release exactly. An output calculated before a rename was registered, or from a year file carrying an out-of-date stamp, is therefore not reused. The error says why, naming the renames the output may lack. Outputs saved before the stamp existed are calculated again once. A stamp or renames record behind a soft link that leads back to itself now counts as unreadable, for year files too, rather than raising. diff --git a/src/policyengine/tax_benefit_models/us/legacy_inputs.py b/src/policyengine/tax_benefit_models/us/legacy_inputs.py index 5607ddf3..0146bebe 100644 --- a/src/policyengine/tax_benefit_models/us/legacy_inputs.py +++ b/src/policyengine/tax_benefit_models/us/legacy_inputs.py @@ -589,7 +589,9 @@ def _read_json_dataset(file: Any, name: str) -> Any: return json.loads(node.asstr()[()]) except (PermissionError, IsADirectoryError, BlockingIOError): raise - except (KeyError, OSError, RecursionError, TypeError, ValueError): + # ``RuntimeError`` covers ``RecursionError``, and h5py raises it when a + # soft link at ``name`` leads back to itself. + except (KeyError, OSError, RuntimeError, TypeError, ValueError): return _UNREADABLE diff --git a/tests/test_us_legacy_inputs_integration.py b/tests/test_us_legacy_inputs_integration.py index 0ff941ab..a2045d2c 100644 --- a/tests/test_us_legacy_inputs_integration.py +++ b/tests/test_us_legacy_inputs_integration.py @@ -1117,6 +1117,47 @@ def test_a_run_over_a_year_file_records_its_stamp(tmp_path): assert _person_outputs(_reloaded(simulation))["wic"].iloc[INFANT] == 0 +def test_a_region_scoped_run_records_its_year_files_stamp(tmp_path): + """A scoped run simulates an in-memory copy of the data, which keeps no + metadata, so the stamp must come from the dataset the run was given.""" + from policyengine.core.scoping_strategy import RowFilterStrategy + + (tmp_path / "stale").mkdir() + (tmp_path / "current").mkdir() + with patch.object(legacy_inputs, "LEGACY_INPUT_RENAMES", {}): + _, cut_before = _create_year_file(tmp_path / "stale") + _, current = _create_year_file(tmp_path / "current") + stale = PolicyEngineUSDataset( + name="stale", description="stale", filepath=cut_before.filepath, year=YEAR + ) + cases = [ + (current, year_file_stamp(), None), + ( + stale, + {"register": {}, "revision": 1}, + "it was calculated from a year file cut before these renames were " + f"registered, which may have lost them: {LEGACY} -> {LIVE}", + ), + ] + for index, (dataset, year_file, problem) in enumerate(cases): + simulation = pe.Simulation( + id=f"scoped-{index}", + dataset=dataset, + tax_benefit_model_version=pe.us.model, + scoping_strategy=RowFilterStrategy( + variable_name="state_code", variable_value="CA" + ), + extra_variables={"person": [LIVE, "wic"]}, + ) + simulation.run() + simulation.save() + path = simulation.output_dataset.filepath + + assert len(_person_outputs(simulation)) == len(PERSON_IDS) + assert read_output_stamp(path) == output_stamp(year_file) + assert output_stamp_problem(path) == problem + + def test_an_output_from_an_out_of_date_year_file_is_not_reused(tmp_path): """A year file cut before the WIC entry lost the draw, so an output calculated from it, under this release's register, is wrong too.""" diff --git a/tests/test_us_output_stamp.py b/tests/test_us_output_stamp.py index 25ae2b37..8073f66e 100644 --- a/tests/test_us_output_stamp.py +++ b/tests/test_us_output_stamp.py @@ -166,6 +166,30 @@ def _write_output(path, *, stamp=_ABSENT, record=_ABSENT, year_file_stamp=_ABSEN return Path(path) +_UNREADABLE = object() + + +def _stored(content): + """The JSON value a stored node holds, decoded as HDF5 stores it. + + ``_ABSENT`` for no node, ``_UNREADABLE`` for a node that does not hold a + UTF-8 JSON scalar, and the value otherwise. This decodes the raw node + kinds of :func:`_write_node` directly, not through the reader under test. + """ + if content is _ABSENT or not isinstance(content, tuple): + return content + kind, value = content + try: + if kind == "text": + return json.loads(value) + if kind == "bytes": + # A fixed-length string reads back without its trailing NUL bytes. + return json.loads(value.rstrip(b"\x00").decode("utf-8")) + except (RecursionError, ValueError): + pass + return _UNREADABLE + + def _same_json(left, right) -> bool: """Equal as JSON text, so ``true`` is not ``1`` and ``1.0`` is not ``1``.""" return json.dumps(left, sort_keys=True) == json.dumps(right, sort_keys=True) @@ -291,22 +315,32 @@ def test_a_missing_or_foreign_file_has_no_output_stamp(tmp_path): # --- Exactness ---------------------------------------------------------- +def _weighted(draw, *choices): + """Draw from one of ``(weight, strategy)`` pairs, in proportion to weight. + + ``st.one_of`` drops a repeated branch, so repeating one does not weight it. + """ + indices = [ + index for index, (weight, _) in enumerate(choices) for _ in range(weight) + ] + return draw(choices[draw(st.sampled_from(indices))][1]) + + @st.composite def near(draw, register): """``register`` itself (most often), one entry away from it, or any.""" - exact = st.just(dict(register)) - return draw( - st.one_of( - exact, - exact, - exact, - REGISTERS, + return _weighted( + draw, + (6, st.just(dict(register))), + (1, REGISTERS), + ( + 1, st.builds( lambda name: {k: v for k, v in register.items() if k != name}, NAMES, ), - st.builds(lambda k, v: {**register, k: v}, NAMES, NAMES), - ) + ), + (1, st.builds(lambda k, v: {**register, k: v}, NAMES, NAMES)), ) @@ -319,28 +353,34 @@ def stamped_outputs(draw): """ register = draw(REGISTERS) stamped = draw(near(register)) - revision = draw(st.sampled_from([REVISION] * 6 + [0, 2, -1, True, 1.0])) - year_file = draw( - st.one_of( - st.none(), - st.none(), + revision = draw(st.sampled_from([REVISION] * 12 + [0, 2, -1, True, 1.0])) + year_file = _weighted( + draw, + (5, st.none()), + ( + 4, st.builds( _year_stamp, near(register), - st.sampled_from([YEAR_FILE_REVISION] * 6 + [0, 2, True, 1.0, "1"]), + st.sampled_from([YEAR_FILE_REVISION] * 12 + [0, 2, True, 1.0, "1"]), ), - st.builds(_year_stamp, near(register)), - st.just("unreadable"), - JSON_VALUES, - ) - ) - subset = st.builds( - lambda keep: {k: v for k, v in stamped.items() if k in keep}, - st.sets(NAMES), + ), + (1, st.just("unreadable")), + (1, JSON_VALUES), ) - record = draw( + record = _weighted( + draw, # A subset of the stamped register, as a run records, most often. - st.one_of(subset, subset, subset, st.just(_ABSENT), REGISTERS, JSON_VALUES) + ( + 6, + st.builds( + lambda keep: {k: v for k, v in stamped.items() if k in keep}, + st.sets(NAMES), + ), + ), + (1, st.just(_ABSENT)), + (1, REGISTERS), + (1, JSON_VALUES), ) return register, _stamp(stamped, revision, year_file), record @@ -422,8 +462,7 @@ def test_the_output_problem_is_total_and_exact_for_any_content(stamp, record, re problem = output_stamp_problem(path) assert problem is None or (isinstance(problem, str) and problem) - if problem is None: - assert _is_current(stamp, record, register) + assert (problem is None) == _is_current(_stored(stamp), _stored(record), register) @pytest.mark.parametrize( @@ -564,6 +603,53 @@ def test_a_stamp_and_record_stored_as_fixed_length_strings_are_read(tmp_path): assert read_output_stamp(path) == output_stamp(year_file_stamp()) +@pytest.mark.parametrize( + "record", + [ + ("text", "{}"), + ("bytes", b"{}"), + ("bytes", b"{}\x00"), + ("text", json.dumps({LEGACY: LIVE})), + ("bytes", json.dumps({LEGACY: LIVE}).encode()), + ], +) +def test_a_record_stored_as_raw_text_or_bytes_is_read(tmp_path, record): + path = _write_output(tmp_path / "output.h5", stamp=output_stamp(), record=record) + + assert output_stamp_problem(path) is None + assert read_renames_record(path) == _stored(record) + + +@pytest.mark.parametrize( + "name", [OUTPUT_STAMP_H5_DATASET, RENAMES_H5_DATASET, YEAR_FILE_STAMP_H5_DATASET] +) +def test_a_soft_link_that_leads_back_to_itself_is_unreadable(tmp_path, name): + """h5py raises ``RuntimeError`` following such a link.""" + path = _write_output( + tmp_path / "output.h5", + stamp=output_stamp(), + record={}, + year_file_stamp=year_file_stamp(), + ) + with h5py.File(path, "a") as file: + del file[name] + file[name] = h5py.SoftLink(f"/{name}") + + if name == OUTPUT_STAMP_H5_DATASET: + assert output_stamp_problem(path) == ( + "its stamp of the register of renamed inputs it was calculated under " + "is malformed" + ) + assert read_output_stamp(path) is None + elif name == RENAMES_H5_DATASET: + assert output_stamp_problem(path) == MALFORMED_RECORD + assert read_renames_record(path) is None + else: + assert year_file_stamp_metadata(path) == "unreadable" + assert read_year_file_stamp(path) is None + assert year_file_stamp_problem(path) is not None + + # --- The year file an output was calculated from ------------------------ From 90fa5037dc28bf7c0b7cf54e1739841558a0dc83 Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Tue, 13 Oct 2026 07:38:08 -0400 Subject: [PATCH 4/6] Apply review 2 nits: exact cyclic year-file reason, reader comment Co-Authored-By: Claude Opus 5.5 --- src/policyengine/tax_benefit_models/us/legacy_inputs.py | 7 ++++--- tests/test_us_output_stamp.py | 4 +++- 2 files changed, 7 insertions(+), 4 deletions(-) diff --git a/src/policyengine/tax_benefit_models/us/legacy_inputs.py b/src/policyengine/tax_benefit_models/us/legacy_inputs.py index 0146bebe..6d17c2d5 100644 --- a/src/policyengine/tax_benefit_models/us/legacy_inputs.py +++ b/src/policyengine/tax_benefit_models/us/legacy_inputs.py @@ -585,12 +585,13 @@ def _read_json_dataset(file: Any, name: str) -> Any: return _UNREADABLE # A numeric dataset raises TypeError; bytes that are not UTF-8, or # not JSON, raise ValueError; JSON nested too deeply raises - # RecursionError. + # RecursionError, a RuntimeError, as does h5py following a soft link + # at ``name`` that leads back to itself. Any RuntimeError counts as + # unreadable content: the errors that say nothing about the content + # are OSErrors, re-raised first. return json.loads(node.asstr()[()]) except (PermissionError, IsADirectoryError, BlockingIOError): raise - # ``RuntimeError`` covers ``RecursionError``, and h5py raises it when a - # soft link at ``name`` leads back to itself. except (KeyError, OSError, RuntimeError, TypeError, ValueError): return _UNREADABLE diff --git a/tests/test_us_output_stamp.py b/tests/test_us_output_stamp.py index 8073f66e..a5b7ab69 100644 --- a/tests/test_us_output_stamp.py +++ b/tests/test_us_output_stamp.py @@ -647,7 +647,9 @@ def test_a_soft_link_that_leads_back_to_itself_is_unreadable(tmp_path, name): else: assert year_file_stamp_metadata(path) == "unreadable" assert read_year_file_stamp(path) is None - assert year_file_stamp_problem(path) is not None + assert year_file_stamp_problem(path) == ( + "its stamp of the register of renamed inputs it was cut under is malformed" + ) # --- The year file an output was calculated from ------------------------ From a51f31c022b4e7d2bdc3cc8c1be88c9c821a9cc3 Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Tue, 13 Oct 2026 07:42:00 -0400 Subject: [PATCH 5/6] Pass save()'s output path as str so the PR adds no mypy errors Co-Authored-By: Claude Opus 5.5 --- src/policyengine/tax_benefit_models/common/model_version.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/src/policyengine/tax_benefit_models/common/model_version.py b/src/policyengine/tax_benefit_models/common/model_version.py index 5cf6a581..cde43087 100644 --- a/src/policyengine/tax_benefit_models/common/model_version.py +++ b/src/policyengine/tax_benefit_models/common/model_version.py @@ -416,8 +416,9 @@ def save(self, simulation: Simulation) -> None: # Stamp last. ``PolicyEngineUSDataset.save()`` truncates the file # first, so an output whose save stopped part-way has no stamp, - # and ``load()`` refuses it. - write_output_stamp(simulation.output_dataset.filepath, stamp) + # and ``load()`` refuses it. It also refuses to save without a + # filepath, so there is one here. + write_output_stamp(str(simulation.output_dataset.filepath), stamp) def load(self, simulation: Simulation) -> None: """Rehydrate the simulation's output dataset from disk. From f06017bcbd58ed8bff71b22907a87dfbf32b52fa Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Tue, 13 Oct 2026 07:46:11 -0400 Subject: [PATCH 6/6] Say which OSErrors the H5 JSON reader re-raises Co-Authored-By: Claude Opus 5.5 --- src/policyengine/tax_benefit_models/us/legacy_inputs.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/policyengine/tax_benefit_models/us/legacy_inputs.py b/src/policyengine/tax_benefit_models/us/legacy_inputs.py index 6d17c2d5..25874f33 100644 --- a/src/policyengine/tax_benefit_models/us/legacy_inputs.py +++ b/src/policyengine/tax_benefit_models/us/legacy_inputs.py @@ -587,8 +587,8 @@ def _read_json_dataset(file: Any, name: str) -> Any: # not JSON, raise ValueError; JSON nested too deeply raises # RecursionError, a RuntimeError, as does h5py following a soft link # at ``name`` that leads back to itself. Any RuntimeError counts as - # unreadable content: the errors that say nothing about the content - # are OSErrors, re-raised first. + # unreadable content. The OSErrors that say nothing about the content + # (no permission, a directory, a lock) are re-raised first. return json.loads(node.asstr()[()]) except (PermissionError, IsADirectoryError, BlockingIOError): raise