diff --git a/changelog.d/532.changed.md b/changelog.d/532.changed.md new file mode 100644 index 00000000..d1a26d1a --- /dev/null +++ b/changelog.d/532.changed.md @@ -0,0 +1 @@ +`create_datasets` stamps each US year file with the register of renamed stored inputs it was cut under and a year-file format revision. `ensure_datasets` creates a year file again, and `load_datasets` refuses it, unless its stamp matches this release exactly. A year file cut before a rename was registered, which may have lost that input, is therefore not reused. `load_datasets` says why it refuses a file, naming the renames the file may lack when its register is the reason. Year files written before the stamp existed are created again once. diff --git a/docs/microsim.md b/docs/microsim.md index 3e1b2da5..1e63432a 100644 --- a/docs/microsim.md +++ b/docs/microsim.md @@ -401,15 +401,49 @@ 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. -Files written before this mapping existed have no record, and are not reused: - -- A saved US output may have been calculated without the mapped inputs. - `load()` refuses it, and `ensure()` runs the simulation again and saves the new - output. -- A year file that `ensure_datasets` or `create_datasets` wrote stores - neither name, so the draw is lost from it. `ensure_datasets` creates such - year files again, and `load_datasets` refuses them. A year file opened - directly, as `PolicyEngineUSDataset(filepath=...)`, is not checked. +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 + +`create_datasets` stamps each year file, after writing everything else, with +the register it was cut under and a year-file format revision +(`legacy_inputs.YEAR_FILE_STAMP_H5_DATASET`, UTF-8 JSON +`{"register": {old: new, ...}, "revision": 1}`). `ensure_datasets` reuses a +year file only when its stamp equals this release's exactly and its renames +record lists only renames in the stamped register. Otherwise it logs why and +creates the year files again; `load_datasets` refuses the file with the same +reason, which names the renames the file may lack when its register is the +cause. To decide, both read only the stamp and the record, never the tables +or the source dataset. `legacy_inputs.year_file_stamp_problem(path)` returns +the reason, or `None` for a current year file. + +A year file is out of date when: + +- it has no stamp. Year files that policyengine.py 6.1.1 and earlier cut, + with policyengine-us 2.x, from a source storing the draw only as + `would_claim_wic` (such as the certified default) store neither WIC take-up + name, so the draw is lost from them. Year files written with the renames + record but before the stamp are cut again once, and so is a year file whose + write did not finish; +- its register lacks an entry this release has. It may have lost that + entry's input, as the older files lost the WIC draw; +- its register has an entry this release lacks, or maps an old name to + another new name. It may store a new name this release's engine does not + define, and a run skips any column the engine does not define; +- it has another format revision, or its stamp or record cannot be read. + +So any change to `LEGACY_INPUT_RENAMES` makes every year file out of date, +and two policyengine.py versions with different registers that share a data +folder cut each other's year files again. `ensure_datasets` raises, rather +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. ## Pinned model versions diff --git a/src/policyengine/tax_benefit_models/us/datasets.py b/src/policyengine/tax_benefit_models/us/datasets.py index 111f4880..a5fbfb9b 100644 --- a/src/policyengine/tax_benefit_models/us/datasets.py +++ b/src/policyengine/tax_benefit_models/us/datasets.py @@ -1,6 +1,7 @@ import hashlib import importlib.util import json +import logging import warnings from importlib import metadata as importlib_metadata from pathlib import Path @@ -31,9 +32,14 @@ pending_legacy_input_renames, read_renames_record, write_renames_record, + write_year_file_stamp, + year_file_stamp, + year_file_stamp_problem, ) from policyengine.utils.hashing import sha256_file +logger = logging.getLogger(__name__) + class USYearData(YearData): """Entity-level data for a single year.""" @@ -371,6 +377,12 @@ def create_datasets( ) -> dict[str, PolicyEngineUSDataset]: """Create PolicyEngineUSDataset instances from logical dataset names or URLs. + Each year file is written to ``{data_folder}/{dataset_stem}_year_{year}.h5`` + and stamped last with the register of renamed stored inputs it was cut + under and the year-file format revision (see + ``legacy_inputs.YEAR_FILE_STAMP_H5_DATASET``), which ``ensure_datasets`` + and ``load_datasets`` check. + Args: datasets: List of logical dataset names or HuggingFace dataset URLs years: List of years to extract data for @@ -394,6 +406,8 @@ def create_datasets( ) dataset_stem = source.name sim = Microsimulation(dataset=source.path, spm=resolve_spm_selection()) + # The register the mapping below uses, stamped into each year file. + stamp = year_file_stamp() # Map stored inputs the engine has since renamed before extracting, # so each year's file stores the live input rather than dropping it. legacy_input_renames = apply_legacy_input_renames_to_microsimulation(sim) @@ -555,6 +569,9 @@ def create_datasets( ), ) us_dataset.save() + # 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) dataset_key = f"{dataset_stem}_{year}" result[dataset_key] = us_dataset @@ -562,16 +579,6 @@ def create_datasets( return result -def _year_file_records_renames(path: Path) -> bool: - """Return whether a year file records the renamed stored inputs mapped. - - ``create_datasets`` writes the record into every year file (``{}`` when - nothing needed mapping). Files written before it did may have lost a - renamed input such as the WIC take-up draw. - """ - return read_renames_record(path) is not None - - def load_datasets( datasets: Optional[list[str]] = None, years: list[int] = [2024, 2025, 2026, 2027, 2028], @@ -579,10 +586,13 @@ def load_datasets( ) -> dict[str, PolicyEngineUSDataset]: """Load PolicyEngineUSDataset instances from saved HDF5 files. - A year file without the record of renamed stored inputs that - ``create_datasets`` writes was cut before those inputs were mapped, and - may have lost one (see ``legacy_inputs.RENAMES_H5_DATASET``), so it is - refused. + A year file is refused unless ``create_datasets`` stamped it with this + release's register of renamed stored inputs and year-file format revision + (see ``legacy_inputs.year_file_stamp_problem``). A file cut before a + register entry was added may have lost that input, as files that + policyengine.py 6.1.1 and earlier cut from the certified default lost the + WIC take-up draw. The error says why, and names the renames the file may + lack when its register is the reason. Args: datasets: List of HuggingFace dataset paths (used to derive file names) @@ -599,17 +609,14 @@ def load_datasets( dataset_stem = dataset_logical_name(resolved_dataset) for year in years: filepath = f"{data_folder}/{dataset_stem}_year_{year}.h5" - if Path(filepath).exists() and not _year_file_records_renames( - Path(filepath) - ): - raise ValueError( - f"US year file {filepath} has no record of the renamed " - "stored inputs mapped when it was cut, so it was written " - "before policyengine.py mapped them and may have lost " - "the WIC take-up draw (every WIC-eligible person would " - "then take WIC up). Regenerate it with ensure_datasets() " - "or create_datasets()." - ) + if Path(filepath).exists(): + problem = year_file_stamp_problem(filepath) + if problem is not None: + raise ValueError( + f"US year file {filepath} is out of date: {problem}. " + "Regenerate it with ensure_datasets() or " + "create_datasets()." + ) us_dataset = PolicyEngineUSDataset( name=f"{dataset_stem}-year-{year}", description=f"US Dataset for year {year} based on {dataset_stem}", @@ -1255,10 +1262,13 @@ def ensure_datasets( ) -> dict[str, PolicyEngineUSDataset]: """Ensure datasets exist, loading if available or creating if not. - Year files without the record of renamed stored inputs that - ``create_datasets`` writes were cut before those inputs were mapped, and - may have lost one such as the WIC take-up draw, so they are created - again rather than loaded. + Existing year files are reused only when ``create_datasets`` stamped them + with this release's register of renamed stored inputs and year-file format + revision (see ``legacy_inputs.year_file_stamp_problem``). Only each file's + stamp and renames record are read to decide, never its tables or the + source dataset. If any year file is missing or out of date, every + year file is created again from the source dataset, and each out-of-date + file is logged with the reason. Args: datasets: List of HuggingFace dataset paths @@ -1270,22 +1280,29 @@ def ensure_datasets( """ datasets = datasets or [get_release_manifest("us").default_dataset] - # Check if all dataset files exist - all_exist = True + # Check that every year file exists and is current. A year file cut + # under another register of renamed stored inputs may have lost one, so + # it is created again rather than reused. + all_current = True for dataset in datasets: resolved_dataset = resolve_dataset_reference("us", dataset) dataset_stem = dataset_logical_name(resolved_dataset) for year in years: filepath = Path(f"{data_folder}/{dataset_stem}_year_{year}.h5") - # A year file written before renamed stored inputs were mapped - # may have lost them, so it is regenerated rather than reused. - if not filepath.exists() or not _year_file_records_renames(filepath): - all_exist = False - break - if not all_exist: - break - - if all_exist: + if not filepath.exists(): + all_current = False + continue + problem = year_file_stamp_problem(filepath) + if problem is not None: + all_current = False + logger.warning( + "US year file %s is out of date: %s. Creating the year " + "files again.", + filepath, + problem, + ) + + if all_current: return load_datasets(datasets=datasets, years=years, data_folder=data_folder) else: return create_datasets( diff --git a/src/policyengine/tax_benefit_models/us/legacy_inputs.py b/src/policyengine/tax_benefit_models/us/legacy_inputs.py index 3b0a2bfa..38fb0249 100644 --- a/src/policyengine/tax_benefit_models/us/legacy_inputs.py +++ b/src/policyengine/tax_benefit_models/us/legacy_inputs.py @@ -40,6 +40,12 @@ #: US default ``populace-us-2024-spm-20260915`` stores the take-up draw as #: ``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. LEGACY_INPUT_RENAMES: dict[str, str] = { "would_claim_wic": "takes_up_wic_if_eligible", } @@ -59,10 +65,37 @@ #: year file). Files written before this mapping existed have none, so its #: absence marks them as calculated or cut without it: ``Simulation.load()`` #: refuses such an output, ``load_datasets`` refuses such a year file and -#: ``ensure_datasets`` regenerates it. A new register entry would need those -#: checks to tell files written before it apart as well. +#: ``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. RENAMES_H5_DATASET = "policyengine_legacy_input_renames" +#: H5 dataset stamping a ``create_datasets`` year file, as UTF-8 JSON +#: ``{"register": {legacy: live, ...}, "revision": int}``: the whole register +#: the file was cut under and its year-file format revision (see +#: :func:`year_file_stamp`). ``create_datasets`` writes it after the tables and +#: the renames record, so a year file whose write stopped part-way has none. +#: ``ensure_datasets`` cuts a year file again, and ``load_datasets`` refuses +#: it, unless its stamp equals this release's (see +#: :func:`year_file_stamp_problem`). ``PolicyEngineUSDataset.save()`` neither +#: writes nor keeps it, so saving a year file again drops the stamp. +YEAR_FILE_STAMP_H5_DATASET = "policyengine_year_file_stamp" + +#: Format revision of ``create_datasets`` year files, stored in their stamp. +#: Raise it when a change to year-file extraction means that files cut before +#: it must be cut again, for a reason the register does not capture (for +#: example, policyengine.py 6.1.1 and earlier exported only the engine's input +#: variables, so a stored column the engine had renamed was dropped). +#: +#: 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. +YEAR_FILE_STAMP_REVISION = 1 + #: ``{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]] @@ -84,39 +117,224 @@ def pending_legacy_input_renames(variables: Any) -> dict[str, str]: } +#: What :func:`_read_json_dataset` returns for a node a file does not have, and +#: for one that is not a scalar dataset holding UTF-8 JSON. +_ABSENT = object() +_UNREADABLE = object() + + def read_renames_record(path: str | Path) -> dict[str, str] | None: """Return the renames record stored in an H5 file, or ``None``. - ``None`` means the file stores no record (see :data:`RENAMES_H5_DATASET`) - or is not an H5 file. + ``None`` means the file stores no record (see :data:`RENAMES_H5_DATASET`), + stores one that is not a JSON object of strings, or is not an H5 file. """ import h5py try: with h5py.File(path, "r") as file: - if RENAMES_H5_DATASET not in file: - return None - raw = file[RENAMES_H5_DATASET].asstr()[()] + record = _read_json_dataset(file, RENAMES_H5_DATASET) except OSError: return None - return dict(json.loads(raw)) + return dict(record) if _is_str_map(record) else None def write_renames_record(path: str | Path, record: Mapping[str, str]) -> None: """Store a renames record in an H5 file, replacing any it has.""" + _write_json_dataset(path, RENAMES_H5_DATASET, dict(record)) + + +def year_file_stamp() -> dict[str, Any]: + """Return the stamp for a year file cut now. + + It holds the whole register, not only the entries an engine needs mapped, + so checking a stamp never imports policyengine-us. The register is read + when this is called. + """ + return { + "register": dict(sorted(LEGACY_INPUT_RENAMES.items())), + "revision": YEAR_FILE_STAMP_REVISION, + } + + +def read_year_file_stamp(path: str | Path) -> Any | None: + """Return the JSON value of a year file's stamp, well formed or not. + + ``None`` means the file stores no readable stamp (see + :data:`YEAR_FILE_STAMP_H5_DATASET`) or is not an H5 file. + """ + import h5py + + try: + with h5py.File(path, "r") as file: + stamp = _read_json_dataset(file, YEAR_FILE_STAMP_H5_DATASET) + except OSError: + return None + return None if stamp is _ABSENT or stamp is _UNREADABLE else stamp + + +def write_year_file_stamp(path: str | Path, stamp: Mapping[str, Any]) -> None: + """Stamp a year file (see :func:`year_file_stamp`), replacing any stamp.""" + _write_json_dataset(path, YEAR_FILE_STAMP_H5_DATASET, dict(stamp)) + + +def year_file_stamp_problem(path: str | Path) -> str | None: + """Return why a ``create_datasets`` year file is out of date, or ``None``. + + A year file is current only when both of these hold: + + - its stamp equals :func:`year_file_stamp`: the same format revision and + the same register; + - it stores the renames record, listing only renames in that register. + + The register must match exactly. A file cut under a register that lacks + an entry may have lost that entry's input, as year files that + policyengine.py 6.1.1 and earlier cut from the certified default lost the + WIC take-up draw. A file cut under an entry this release does not + register, or registers with another live input, may store a live name + this release's engine does not define, and a run skips any column the + engine does not define. + + Only the stamp and the record are read, never the tables or the source + dataset. A file that cannot be read as HDF5 is out of date. An error that + says nothing about the file's contents (no permission, a directory, a file + another process has locked) is raised instead, since the file could not + be written again either. + + Returns: + ``None`` for a current year file; otherwise a clause, starting "it" + or "its", that says why the file is out of date and names the + renames it may have lost. + """ + import h5py + + current = year_file_stamp() + register = current["register"] + try: + with h5py.File(path, "r") as file: + stamp = _read_json_dataset(file, YEAR_FILE_STAMP_H5_DATASET) + record = _read_json_dataset(file, RENAMES_H5_DATASET) + except (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 != YEAR_FILE_STAMP_REVISION: + return ( + f"it was cut with year-file format revision {revision}, and this " + f"version of policyengine.py reads only revision " + f"{YEAR_FILE_STAMP_REVISION}" + ) + if stamp is _ABSENT: + if record is _ABSENT: + return ( + "it has no record of the renamed stored inputs mapped when it " + "was cut, so it was cut by an older policyengine.py or not " + "written completely" + + ( + f", and may have lost any of: {_describe(register)}" + if register + else "" + ) + ) + return "it has no stamp of the register of renamed inputs it was cut under" + ( + f", so it may have lost any of: {_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"} + and type(revision) is int + and _is_str_map(stamp["register"]) + ): + return ( + "its stamp of the register of renamed inputs it was cut under is malformed" + ) + if record is _ABSENT: + return "it has no record of the renamed stored inputs mapped when it was cut" + if not _is_str_map(record): + return ( + "its record of the renamed stored inputs mapped when it was cut 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)}" + ) + 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 + } + reasons = [] + if missing: + reasons.append( + "it was cut before these renames were registered, so it may have " + f"lost them: {_describe(missing)}" + ) + if unknown: + reasons.append( + "it was cut under renames this version of policyengine.py does not " + f"register: {_describe(unknown)}" + ) + return "; ".join(reasons) + + +def _read_json_dataset(file: Any, name: str) -> Any: + """Return the JSON value stored in an open H5 file's ``name`` dataset.""" + import h5py + + try: + node = file.get(name) + if node is None: + return _ABSENT + if not isinstance(node, h5py.Dataset) or node.shape != (): + return _UNREADABLE + # A numeric dataset raises TypeError; bytes that are not UTF-8, or + # not JSON, raise ValueError; JSON nested too deeply raises + # RecursionError. + return json.loads(node.asstr()[()]) + except (PermissionError, IsADirectoryError, BlockingIOError): + raise + except (KeyError, OSError, RecursionError, TypeError, ValueError): + return _UNREADABLE + + +def _write_json_dataset(path: str | Path, name: str, value: Any) -> None: + """Store ``value`` as UTF-8 JSON in an H5 file's ``name`` dataset.""" import h5py with h5py.File(path, "a") as file: - if RENAMES_H5_DATASET in file: - del file[RENAMES_H5_DATASET] + if name in file: + del file[name] # UTF-8 JSON in a dataset, as for the SPM receipt of a saved output. file.create_dataset( - RENAMES_H5_DATASET, - data=json.dumps(dict(record), sort_keys=True), + name, + data=json.dumps(value, sort_keys=True), dtype=h5py.string_dtype("utf-8"), ) +def _is_str_map(value: Any) -> bool: + return isinstance(value, dict) and all( + isinstance(key, str) and isinstance(item, str) for key, item in value.items() + ) + + +def _describe(renames: Mapping[str, str]) -> str: + return ", ".join(f"{legacy} -> {live}" for legacy, live in sorted(renames.items())) + + def check_yearly_periods(column: str, periods: Iterable[Any], source: Any) -> None: """Refuse a legacy column stored for any period that is not a year. diff --git a/tests/test_us_legacy_inputs_integration.py b/tests/test_us_legacy_inputs_integration.py index 68c1a65e..182d39af 100644 --- a/tests/test_us_legacy_inputs_integration.py +++ b/tests/test_us_legacy_inputs_integration.py @@ -13,6 +13,8 @@ from __future__ import annotations +from unittest.mock import patch + import h5py import numpy as np import pandas as pd @@ -25,6 +27,7 @@ pytest.importorskip("spm_calculator.policyengine_adapter") import policyengine as pe # noqa: E402 +from policyengine.tax_benefit_models.us import datasets as us_datasets # noqa: E402 from policyengine.tax_benefit_models.us import legacy_inputs # noqa: E402 from policyengine.tax_benefit_models.us.datasets import ( # noqa: E402 PolicyEngineUSDataset, @@ -35,9 +38,13 @@ ) from policyengine.tax_benefit_models.us.legacy_inputs import ( # noqa: E402 RENAMES_H5_DATASET, + YEAR_FILE_STAMP_H5_DATASET, apply_legacy_input_renames, apply_legacy_input_renames_to_microsimulation, read_renames_record, + read_year_file_stamp, + year_file_stamp, + year_file_stamp_problem, ) from policyengine.tax_benefit_models.us.model import ( # noqa: E402 PolicyEngineUSLatest, @@ -558,6 +565,9 @@ def test_create_datasets_extracts_the_draw_under_the_live_name(tmp_path): assert dataset.metadata["legacy_input_renames"] == RENAME # The year file keeps the record, so reloading it keeps the chain. assert read_renames_record(dataset.filepath) == RENAME + # It is stamped with the register it was cut under, so it is current. + assert read_year_file_stamp(dataset.filepath) == year_file_stamp() + assert year_file_stamp_problem(dataset.filepath) is None reloaded = PolicyEngineUSDataset( name=dataset.name, description=dataset.description, @@ -632,6 +642,194 @@ def test_ensure_datasets_regenerates_a_year_file_cut_before_the_mapping(tmp_path assert simulation.release_bundle["legacy_input_renames"] == RENAME +# --- The year-file stamp: the register a year file was cut under -------- + + +def _ensure(source, data_folder): + (dataset,) = ensure_datasets( + datasets=[source], + years=[YEAR], + data_folder=str(data_folder), + allow_unmanaged=True, + ).values() + return dataset + + +def _assert_reused(source, data_folder): + """A current year file is loaded, never cut again.""" + + def refuse(*args, **kwargs): + raise AssertionError("a current year file was cut again") + + with patch.object(us_datasets, "create_datasets", refuse): + return _ensure(source, data_folder) + + +def test_a_year_file_cut_before_a_register_entry_is_regenerated(tmp_path): + """Replay the WIC history without a real engine rename. + + Cutting under an empty register is what policyengine.py did before the + WIC entry existed: the engine skips the stored legacy column, so the year + file stores neither name. + """ + with patch.object(legacy_inputs, "LEGACY_INPUT_RENAMES", {}): + source, dataset = _create_year_file(tmp_path) + # Under the register it was cut with, the file is current. + assert year_file_stamp_problem(dataset.filepath) is None + path = dataset.filepath + person = pd.DataFrame(dataset.data.person) + assert LIVE not in person.columns + assert LEGACY not in person.columns + assert read_renames_record(path) == {} + assert read_year_file_stamp(path) == {"register": {}, "revision": 1} + # The draw is lost from it: a run gives the infant WIC. + assert _person_outputs(_run(dataset))["wic"].iloc[INFANT] > 0 + + # Under this release's register, the file is out of date. + reason = ( + "it was cut before these renames were registered, so it may have lost " + f"them: {LEGACY} -> {LIVE}" + ) + assert year_file_stamp_problem(path) == reason + with pytest.raises(ValueError, match=f"is out of date: {reason}"): + load_datasets( + datasets=[source], years=[YEAR], data_folder=str(tmp_path / "data") + ) + + regenerated = _ensure(source, tmp_path / "data") + + assert regenerated.filepath == path + assert read_year_file_stamp(path) == year_file_stamp() + assert pd.DataFrame(regenerated.data.person)[LIVE].tolist() == DRAW + loaded = _assert_reused(source, tmp_path / "data") + assert _person_outputs(_run(loaded))["wic"].iloc[INFANT] == 0 + + +def test_a_future_register_entry_regenerates_year_files(tmp_path): + """An entry the engine does not need still makes older year files stale.""" + from policyengine_us.system import system + + future = {"fictitious_old_input": "fictitious_new_input"} + for name in (*future, *future.values()): + assert name not in system.variables + source, dataset = _create_year_file(tmp_path) + data_folder = tmp_path / "data" + + with patch.object(legacy_inputs, "LEGACY_INPUT_RENAMES", {**RENAME, **future}): + reason = year_file_stamp_problem(dataset.filepath) + assert reason == ( + "it was cut before these renames were registered, so it may have " + "lost them: fictitious_old_input -> fictitious_new_input" + ) + with pytest.raises(ValueError, match="fictitious_old_input"): + load_datasets(datasets=[source], years=[YEAR], data_folder=str(data_folder)) + + created = [] + create = us_datasets.create_datasets + + def spy(*args, **kwargs): + created.append(kwargs["years"]) + return create(*args, **kwargs) + + with patch.object(us_datasets, "create_datasets", spy): + regenerated = _ensure(source, data_folder) + + assert created == [[YEAR]] + assert read_year_file_stamp(dataset.filepath) == { + "register": {**RENAME, **future}, + "revision": 1, + } + # Only the entry the engine needs was applied, and the data is as + # before. + assert read_renames_record(dataset.filepath) == RENAME + assert pd.DataFrame(regenerated.data.person)[LIVE].tolist() == DRAW + _assert_reused(source, data_folder) + + # Back under this release's register, the newer file is out of date too. + assert "does not register: fictitious_old_input" in year_file_stamp_problem( + dataset.filepath + ) + + +def test_ensure_datasets_regenerates_a_year_file_with_only_the_renames_record( + tmp_path, +): + """A year file written with the record but before the stamp existed.""" + source, dataset = _create_year_file(tmp_path) + with h5py.File(dataset.filepath, "a") as file: + del file[YEAR_FILE_STAMP_H5_DATASET] + data_folder = tmp_path / "data" + + with pytest.raises(ValueError, match="it has no stamp of the register"): + load_datasets(datasets=[source], years=[YEAR], data_folder=str(data_folder)) + + regenerated = _ensure(source, data_folder) + + assert read_year_file_stamp(regenerated.filepath) == year_file_stamp() + assert pd.DataFrame(regenerated.data.person)[LIVE].tolist() == DRAW + _assert_reused(source, data_folder) + + +def test_create_datasets_stamps_each_year_file_after_saving_it(tmp_path): + source = _write_entity_tables(tmp_path / "populace_layout.h5") + events = [] + save = PolicyEngineUSDataset.save + stamp = us_datasets.write_year_file_stamp + + def recording_save(self): + save(self) + events.append(("save", self.filepath)) + + def recording_stamp(path, value): + stamp(path, value) + events.append(("stamp", path)) + + with ( + patch.object(PolicyEngineUSDataset, "save", recording_save), + patch.object(us_datasets, "write_year_file_stamp", recording_stamp), + ): + created = create_datasets( + datasets=[source], + years=[YEAR, YEAR + 1], + data_folder=str(tmp_path / "data"), + allow_unmanaged=True, + ) + + paths = [dataset.filepath for dataset in created.values()] + assert events == [ + ("save", paths[0]), + ("stamp", paths[0]), + ("save", paths[1]), + ("stamp", paths[1]), + ] + for path in paths: + assert year_file_stamp_problem(path) is None + + +def test_a_year_file_whose_stamp_was_never_written_is_out_of_date(tmp_path): + source = _write_entity_tables(tmp_path / "populace_layout.h5") + + def stop(path, value): + raise RuntimeError("stopped before stamping") + + with patch.object(us_datasets, "write_year_file_stamp", stop): + with pytest.raises(RuntimeError, match="stopped before stamping"): + create_datasets( + datasets=[source], + years=[YEAR], + data_folder=str(tmp_path / "data"), + allow_unmanaged=True, + ) + + (path,) = (tmp_path / "data").glob("*_year_*.h5") + # The tables and the record were written, but without the stamp the file + # is cut again. + assert read_renames_record(path) == RENAME + assert year_file_stamp_problem(path).startswith("it has no stamp") + assert _ensure(source, tmp_path / "data").filepath == str(path) + assert year_file_stamp_problem(path) is None + + # --- Saved outputs and run records keep the record ---------------------- diff --git a/tests/test_us_year_file_stamp.py b/tests/test_us_year_file_stamp.py new file mode 100644 index 00000000..7d11eef0 --- /dev/null +++ b/tests/test_us_year_file_stamp.py @@ -0,0 +1,1012 @@ +"""Unit and property tests for the stamp on US year files. + +``create_datasets`` stamps each year file with the register of renamed stored +inputs it was cut under and a year-file format revision +(``legacy_inputs.YEAR_FILE_STAMP_H5_DATASET``). ``ensure_datasets`` cuts a +year file again, and ``load_datasets`` refuses it, 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.** A year file is current exactly when its stamp has this + release's format revision and register, and its renames record lists only + renames in that 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 (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 file's register and + this release's disagree on, from the register alone. +- **No partial file is current.** Every truncation of a current year file, + every save that stops part-way and every save that does not stamp is out + of date. +- **Ensure agrees with load.** ``ensure_datasets`` cuts the year files again + exactly when one is missing or ``load_datasets`` would refuse one, and + decides from the stamps alone. + +``test_us_legacy_inputs_integration.py`` runs ``create_datasets`` and +``ensure_datasets`` 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 MagicMock, patch + +import h5py +import numpy as np +import pandas as pd +import pytest +from hypothesis import given, settings +from hypothesis import strategies as st +from microdf import MicroDataFrame + +from policyengine.tax_benefit_models.us import datasets as us_datasets +from policyengine.tax_benefit_models.us import legacy_inputs +from policyengine.tax_benefit_models.us.datasets import ( + PolicyEngineUSDataset, + USYearData, + _is_policyengine_core_h5, + ensure_datasets, + load_datasets, +) +from policyengine.tax_benefit_models.us.legacy_inputs import ( + LEGACY_INPUT_RENAMES, + RENAMES_H5_DATASET, + YEAR_FILE_STAMP_H5_DATASET, + YEAR_FILE_STAMP_REVISION, + read_renames_record, + read_year_file_stamp, + write_renames_record, + write_year_file_stamp, + year_file_stamp, + year_file_stamp_problem, +) + +LEGACY = "would_claim_wic" +LIVE = "takes_up_wic_if_eligible" +REVISION = YEAR_FILE_STAMP_REVISION +#: ``hf://`` references are used as they are, and name year files after the +#: file's stem, so no manifest or download is involved. +SOURCE = "hf://fixture/tiny.h5" +STEM = "tiny" +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, +) + + +def _frames() -> dict[str, pd.DataFrame]: + person = pd.DataFrame( + { + "person_id": [1, 2], + **{f"person_{group}_id": [1, 1] for group in GROUPS}, + "person_weight": [1.0, 1.0], + LIVE: [True, False], + } + ) + frames = {"person": person} + for group in GROUPS: + frames[group] = pd.DataFrame({f"{group}_id": [1], f"{group}_weight": [1.0]}) + return frames + + +def _year_dataset(path) -> PolicyEngineUSDataset: + return PolicyEngineUSDataset( + name="year-file", + description="Uncertified two-person year-file 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_year_file(path, *, stamp=_ABSENT, record=_ABSENT): + """Write a year file's tables, then its record and stamp as given. + + ``stamp`` and ``record`` are JSON values, ``_ABSENT``, or a raw node + (see :func:`_write_node`). + """ + Path(path).parent.mkdir(parents=True, exist_ok=True) + _year_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, stamp) + return Path(path) + + +def _is_current(stamp, record, register) -> bool: + """The exactness invariant, stated directly on the stored JSON values.""" + + def str_map(value): + return isinstance(value, dict) and all( + isinstance(key, str) and isinstance(item, str) + for key, item in value.items() + ) + + return ( + isinstance(stamp, dict) + and set(stamp) == {"register", "revision"} + and type(stamp["revision"]) is int + and stamp["revision"] == REVISION + and str_map(stamp["register"]) + and stamp["register"] == register + and str_map(record) + and record.items() <= stamp["register"].items() + ) + + +def _stamp(register, revision=REVISION): + return {"register": dict(register), "revision": revision} + + +# --- The stamp ---------------------------------------------------------- + + +def test_the_stamp_holds_the_whole_register_and_the_format_revision(): + assert year_file_stamp() == { + "register": {LEGACY: LIVE}, + "revision": 1, + } + assert year_file_stamp()["register"] == LEGACY_INPUT_RENAMES + + +def test_the_stamp_reads_the_register_when_called(monkeypatch): + """datasets.py must not freeze the register at import.""" + monkeypatch.setattr(legacy_inputs, "LEGACY_INPUT_RENAMES", {"b": "c", "a": "d"}) + + assert year_file_stamp() == {"register": {"a": "d", "b": "c"}, "revision": 1} + + +def test_the_stamp_is_stored_as_utf8_json(tmp_path): + path = _write_year_file(tmp_path / "year.h5") + + write_year_file_stamp(path, year_file_stamp()) + + with h5py.File(path, "r") as file: + assert file[YEAR_FILE_STAMP_H5_DATASET].asstr()[()] == ( + f'{{"register": {{"{LEGACY}": "{LIVE}"}}, "revision": 1}}' + ) + + +@settings(max_examples=50, deadline=None) +@given( + first=st.fixed_dictionaries({"register": REGISTERS, "revision": st.integers()}), + second=st.fixed_dictionaries({"register": REGISTERS, "revision": st.integers()}), +) +def test_a_year_file_stamp_reads_back_as_written(first, second): + """Round trip: a stamp reads back exactly, a new one replaces it, and + the tables and record are left alone.""" + with tempfile.TemporaryDirectory() as directory: + path = _write_year_file(Path(directory) / "year.h5", record={}) + assert read_year_file_stamp(path) is None + + write_year_file_stamp(path, first) + assert read_year_file_stamp(path) == first + + write_year_file_stamp(path, second) + assert read_year_file_stamp(path) == second + assert read_renames_record(path) == {} + with pd.HDFStore(path, mode="r") as store: + assert store["person"]["person_id"].tolist() == [1, 2] + + +def test_a_missing_or_foreign_file_has_no_stamp(tmp_path): + assert read_year_file_stamp(tmp_path / "missing.h5") is None + text = tmp_path / "notes.txt" + text.write_text("not an H5 file") + assert read_year_file_stamp(text) is None + + +def test_the_stamp_does_not_disturb_pandas_or_the_layout_check(tmp_path): + path = _write_year_file(tmp_path / "year.h5", stamp=year_file_stamp(), record={}) + + with pd.HDFStore(path, mode="r") as store: + assert sorted(store.keys()) == sorted(f"/{key}" for key in _frames()) + assert not _is_policyengine_core_h5(path) + loaded = PolicyEngineUSDataset( + name="year-file", description="reload", filepath=str(path), year=2024 + ) + assert pd.DataFrame(loaded.data.person)[LIVE].tolist() == [True, False] + assert loaded.metadata["legacy_input_renames"] == {} + + +# --- Exactness ---------------------------------------------------------- + + +@st.composite +def stamped_files(draw): + """A current register, and a year file's stamp and record near it.""" + register = draw(REGISTERS) + stamped = draw( + st.one_of( + st.just(dict(register)), + REGISTERS, + # One entry dropped, added or pointed elsewhere. + 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), + ) + ) + revision = draw(st.sampled_from([REVISION, REVISION, 0, 2, -1, 99])) + record = draw( + st.one_of( + st.just(_ABSENT), + # A subset of the stamped register, as create_datasets writes. + st.builds( + lambda keep: {k: v for k, v in stamped.items() if k in keep}, + st.sets(NAMES), + ), + REGISTERS, + ) + ) + return register, _stamp(stamped, revision), record + + +@settings(max_examples=300, deadline=None) +@given(case=stamped_files()) +def test_a_year_file_is_current_only_under_this_releases_exact_stamp(case): + """Exactness, and refinement of the record check a year file had before. + + A subset rule would accept a file cut by a newer release under a larger + register, which stores live names this release's engine may not define. + """ + register, stamp, record = case + with tempfile.TemporaryDirectory() as directory: + path = _write_year_file(Path(directory) / "year.h5", stamp=stamp, record=record) + with patch.object(legacy_inputs, "LEGACY_INPUT_RENAMES", register): + problem = year_file_stamp_problem(path) + + assert (problem is None) == _is_current(stamp, record, register) + assert problem is None or (isinstance(problem, str) and problem) + if problem is None: + # Every file the stamp accepts, the record check accepted too. + assert read_renames_record(path) == record + + +#: 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, +) + +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 + ), + } + ), + 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_year_file_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_year_file(Path(directory) / "year.h5", stamp=stamp, record=record) + with patch.object(legacy_inputs, "LEGACY_INPUT_RENAMES", register): + problem = year_file_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}, + {"register": {LEGACY: LIVE}, "revision": 1.0}, + {"register": {LEGACY: LIVE}, "revision": "1"}, + {"register": {LEGACY: LIVE}}, + {"register": {LEGACY: LIVE}, "revision": 1, "extra": None}, + {"register": {LEGACY: 5}, "revision": 1}, + {"register": [[LEGACY, LIVE]], "revision": 1}, + [{"register": {LEGACY: LIVE}, "revision": 1}], + ("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_year_file(tmp_path / "year.h5", stamp=stamp, record={}) + + assert year_file_stamp_problem(path) == ( + "its stamp of the register of renamed inputs it was cut under is malformed" + ) + + +@pytest.mark.parametrize("revision", [0, 2, -1]) +def test_a_year_file_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_year_file(tmp_path / "year.h5", stamp=stamp) + + assert year_file_stamp_problem(path) == ( + f"it was cut with year-file format revision {revision}, and this " + "version of policyengine.py reads only revision 1" + ) + + +def test_files_without_a_stamp_are_out_of_date_with_distinct_reasons(tmp_path): + """policyengine.py 6.1.1 and earlier wrote neither; #531 wrote the record.""" + oldest = _write_year_file(tmp_path / "oldest.h5") + record_only = _write_year_file(tmp_path / "record_only.h5", record={LEGACY: LIVE}) + + assert year_file_stamp_problem(oldest) == ( + "it has no record of the renamed stored inputs mapped when it was cut, " + "so it was cut by an older policyengine.py or not written completely, " + f"and may have lost any of: {LEGACY} -> {LIVE}" + ) + assert year_file_stamp_problem(record_only) == ( + "it has no stamp of the register of renamed inputs it was cut under, so " + f"it may have lost any of: {LEGACY} -> {LIVE}" + ) + + +def test_an_unstamped_file_is_out_of_date_even_with_an_empty_register( + tmp_path, monkeypatch +): + monkeypatch.setattr(legacy_inputs, "LEGACY_INPUT_RENAMES", {}) + oldest = _write_year_file(tmp_path / "oldest.h5") + record_only = _write_year_file(tmp_path / "record_only.h5", record={}) + + assert year_file_stamp_problem(oldest) == ( + "it has no record of the renamed stored inputs mapped when it was cut, " + "so it was cut by an older policyengine.py or not written completely" + ) + assert year_file_stamp_problem(record_only) == ( + "it has no stamp of the register of renamed inputs it was cut under" + ) + + +@pytest.mark.parametrize( + "record, reason", + [ + ( + _ABSENT, + "it has no record of the renamed stored inputs mapped when it was cut", + ), + ( + [], + "its record of the renamed stored inputs mapped when it was cut is malformed", + ), + ( + [[LEGACY, LIVE]], + "its record of the renamed stored inputs mapped when it was cut is malformed", + ), + ( + 5, + "its record of the renamed stored inputs mapped when it was cut is malformed", + ), + ( + {LEGACY: 1}, + "its record of the renamed stored inputs mapped when it was cut is malformed", + ), + ( + ("text", "{"), + "its record of the renamed stored inputs mapped when it was cut is malformed", + ), + ( + ("group", None), + "its record of the renamed stored inputs mapped when it was cut is malformed", + ), + ( + {"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_file_needs_a_record_within_its_register(tmp_path, record, reason): + path = _write_year_file( + tmp_path / "year.h5", stamp=year_file_stamp(), record=record + ) + + assert year_file_stamp_problem(path) == reason + + +def test_a_stamp_and_record_stored_as_fixed_length_strings_are_read(tmp_path): + stamp = json.dumps(year_file_stamp()).encode() + record = json.dumps({LEGACY: LIVE}).encode() + path = _write_year_file( + tmp_path / "year.h5", stamp=("bytes", stamp), record=("bytes", record) + ) + + assert year_file_stamp_problem(path) is None + assert read_year_file_stamp(path) == year_file_stamp() + assert read_renames_record(path) == {LEGACY: LIVE} + + +def test_raising_the_format_revision_makes_older_files_out_of_date( + tmp_path, monkeypatch +): + older = _write_year_file(tmp_path / "older.h5", stamp=year_file_stamp(), record={}) + monkeypatch.setattr(legacy_inputs, "YEAR_FILE_STAMP_REVISION", 7) + newer = _write_year_file(tmp_path / "newer.h5", stamp=year_file_stamp(), record={}) + + assert year_file_stamp()["revision"] == 7 + assert year_file_stamp_problem(newer) is None + assert year_file_stamp_problem(older) == ( + "it was cut with year-file format revision 1, and this version of " + "policyengine.py reads only revision 7" + ) + + +#: JSON nested very deeply. Whether the parser reaches its limit depends on +#: the interpreter and the platform: CPython 3.14 checks the native stack, so +#: it varies with stack size, architecture and build. In our runs, CPython +#: 3.11-3.13, and 3.14 on macOS, raised RecursionError, while 3.14 on CI's +#: Linux runners parsed it into nested lists. Either way the stamp or record +#: is not well formed. +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 cut under is malformed", + ), + ( + year_file_stamp(), + ("text", DEEP_JSON), + "its record of the renamed stored inputs mapped when it was cut is " + "malformed", + ), + ( + {"revision": 2}, + ("text", DEEP_JSON), + "it was cut with year-file format revision 2, and this version of " + "policyengine.py reads only revision 1", + ), + ], +) +def test_json_nested_too_deeply_is_out_of_date_not_an_error( + tmp_path, stamp, record, reason +): + path = _write_year_file(tmp_path / "year.h5", stamp=stamp, record=record) + + assert year_file_stamp_problem(path) == reason + if isinstance(stamp, tuple): + # Unreadable (None) or parsed into nested lists, but never a stamp. + assert not isinstance(read_year_file_stamp(path), dict) + else: + assert read_year_file_stamp(path) == stamp + assert read_renames_record(path) == (None if isinstance(record, tuple) else record) + + +@pytest.mark.parametrize("name", [YEAR_FILE_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 (see ``DEEP_JSON``).""" + path = _write_year_file(tmp_path / "year.h5", stamp=year_file_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 year_file_stamp_problem(path) == ( + "its stamp of the register of renamed inputs it was cut under is malformed" + if name == YEAR_FILE_STAMP_H5_DATASET + else "its record of the renamed stored inputs mapped when it was cut is " + "malformed" + ) + + +def test_dangling_links_in_place_of_the_stamp_and_record_read_as_absent(tmp_path): + path = _write_year_file(tmp_path / "year.h5") + with h5py.File(path, "a") as file: + file[YEAR_FILE_STAMP_H5_DATASET] = h5py.SoftLink("/nowhere") + file[RENAMES_H5_DATASET] = h5py.ExternalLink( + str(tmp_path / "missing.h5"), "/record" + ) + + assert year_file_stamp_problem(path).startswith( + "it has no record of the renamed stored inputs mapped when it was cut" + ) + assert read_year_file_stamp(path) is None + assert read_renames_record(path) is None + + +@pytest.mark.parametrize("record", [{}, {LEGACY: LIVE}]) +def test_a_file_stamped_with_this_releases_register_is_current(tmp_path, record): + path = _write_year_file( + tmp_path / "year.h5", stamp=year_file_stamp(), record=record + ) + + assert year_file_stamp_problem(path) is None + + +# --- Named renames ------------------------------------------------------ + + +@settings(max_examples=200, deadline=None) +@given(register=REGISTERS, stamped=REGISTERS) +def test_a_stale_year_file_names_each_rename_it_may_lack(register, stamped): + """The reason names every entry the registers disagree on, as + ``old -> new`` from the register that has it, and no other.""" + with tempfile.TemporaryDirectory() as directory: + path = _write_year_file( + Path(directory) / "year.h5", stamp=_stamp(stamped), record={} + ) + with patch.object(legacy_inputs, "LEGACY_INPUT_RENAMES", register): + problem = year_file_stamp_problem(path) + + if stamped == register: + assert problem is None + return + 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} + for legacy, live in missing.items(): + assert f"{legacy} -> {live}" in problem.split("; ")[0] + for legacy, live in unknown.items(): + assert f"{legacy} -> {live}" in problem.split("; ")[-1] + assert ("cut before these renames were registered" in problem) == bool(missing) + assert ("does not register" in problem) == bool(unknown) + for legacy, live in register.items() & stamped.items(): + assert f"{legacy} -> {live}" not in problem + + +def test_a_future_register_entry_makes_a_stamped_file_out_of_date( + tmp_path, monkeypatch +): + path = _write_year_file(tmp_path / "year.h5", stamp=year_file_stamp(), record={}) + monkeypatch.setattr( + legacy_inputs, + "LEGACY_INPUT_RENAMES", + {**LEGACY_INPUT_RENAMES, "old_input": "new_input"}, + ) + + assert year_file_stamp_problem(path) == ( + "it was cut before these renames were registered, so it may have lost " + "them: old_input -> new_input" + ) + + +@pytest.mark.parametrize( + "register, reason", + [ + ( + {}, + "it was cut under renames this version of policyengine.py does not " + f"register: {LEGACY} -> {LIVE}", + ), + ( + {LEGACY: "elsewhere"}, + "it was cut before these renames were registered, so it may have lost " + f"them: {LEGACY} -> elsewhere; it was cut under renames this version " + f"of policyengine.py does not register: {LEGACY} -> {LIVE}", + ), + ], +) +def test_retiring_or_retargeting_an_entry_makes_a_stamped_file_out_of_date( + tmp_path, monkeypatch, register, reason +): + path = _write_year_file( + tmp_path / "year.h5", stamp=year_file_stamp(), record={LEGACY: LIVE} + ) + monkeypatch.setattr(legacy_inputs, "LEGACY_INPUT_RENAMES", register) + + # The record of a file cut under the old register lies within that + # register, so the register is what makes it out of date. + assert year_file_stamp_problem(path) == reason + + +# --- No partial file is current ----------------------------------------- + + +@functools.cache +def _current_year_file_bytes() -> bytes: + with tempfile.TemporaryDirectory() as directory: + path = _write_year_file( + Path(directory) / "year.h5", stamp=year_file_stamp(), record={} + ) + assert year_file_stamp_problem(path) is None + return path.read_bytes() + + +@settings(max_examples=100, deadline=None) +@given(data=st.data()) +def test_a_truncated_year_file_is_never_current(data): + content = _current_year_file_bytes() + length = data.draw(st.integers(0, len(content) - 1)) + with tempfile.TemporaryDirectory() as directory: + path = Path(directory) / "year.h5" + path.write_bytes(content[:length]) + + assert year_file_stamp_problem(path) is not None + + +def test_a_file_that_is_not_hdf5_is_out_of_date(tmp_path): + path = tmp_path / "year.h5" + path.write_text("not an H5 file") + + assert year_file_stamp_problem(path) == ( + "it cannot be read as an HDF5 file (it may be half-written or corrupt)" + ) + + +def test_saving_a_year_file_again_drops_its_stamp(tmp_path): + path = _write_year_file(tmp_path / "year.h5", stamp=year_file_stamp(), record={}) + + dataset = PolicyEngineUSDataset( + name="year-file", description="reload", filepath=str(path), year=2024 + ) + dataset.save() + + # ``save()`` rewrites the tables and the record it loaded, but not the + # stamp, which only ``create_datasets`` writes. + assert read_renames_record(path) == {} + assert read_year_file_stamp(path) is None + assert year_file_stamp_problem(path).startswith("it has no stamp") + + +def test_a_save_that_stops_part_way_leaves_no_stamp(tmp_path, monkeypatch): + path = _write_year_file(tmp_path / "year.h5", stamp=year_file_stamp(), record={}) + put = pd.HDFStore.put + + def put_until_family(self, key, value, *args, **kwargs): + if key == "family": + raise RuntimeError("stopped part-way") + return put(self, key, value, *args, **kwargs) + + monkeypatch.setattr(pd.HDFStore, "put", put_until_family) + with pytest.raises(RuntimeError, match="stopped part-way"): + _year_dataset(path).save() + + assert read_year_file_stamp(path) is None + assert year_file_stamp_problem(path) is not None + + +# --- Errors unrelated to the file's contents ---------------------------- + + +def _unreadable_by_permission(path) -> Path: + path = _write_year_file(path, stamp=year_file_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") + data_folder = tmp_path / "data" + data_folder.mkdir() + path = make(data_folder / f"{STEM}_year_2024.h5") + create = MagicMock() + try: + with pytest.raises(error): + year_file_stamp_problem(path) + with pytest.raises(error): + load_datasets(datasets=[SOURCE], years=[2024], data_folder=str(data_folder)) + with patch.object(us_datasets, "create_datasets", create): + with pytest.raises(error): + ensure_datasets( + datasets=[SOURCE], years=[2024], data_folder=str(data_folder) + ) + finally: + if path.is_file(): + os.chmod(path, 0o644) + create.assert_not_called() + + +def test_a_locked_file_is_raised_not_cut_again(tmp_path, monkeypatch): + path = _write_year_file(tmp_path / "year.h5", stamp=year_file_stamp(), record={}) + + def locked(*args, **kwargs): + raise BlockingIOError(35, "Unable to lock file") + + monkeypatch.setattr(h5py, "File", locked) + with pytest.raises(BlockingIOError): + year_file_stamp_problem(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_year_file(tmp_path / "year.h5", stamp=year_file_stamp(), record={}) + + def fail(self, *args, **kwargs): + raise error + + monkeypatch.setattr(h5py.Dataset, "asstr", fail) + with pytest.raises(type(error)): + year_file_stamp_problem(path) + + +# --- read_renames_record reads only well-formed records ----------------- + + +@pytest.mark.parametrize( + "record", + [[], [[LEGACY, LIVE]], 5, {LEGACY: 1}, ("text", "{"), ("group", None)], +) +def test_a_malformed_renames_record_reads_as_none(tmp_path, record): + path = _write_year_file(tmp_path / "year.h5", record=record) + + assert read_renames_record(path) is None + + +def test_a_well_formed_renames_record_still_reads_back(tmp_path): + path = _write_year_file(tmp_path / "year.h5") + write_renames_record(path, {LEGACY: LIVE}) + + assert read_renames_record(path) == {LEGACY: LIVE} + + +# --- ensure_datasets agrees with load_datasets -------------------------- + +#: Each year file state, and whether it is current. +STATES = { + "current": (dict(stamp=year_file_stamp(), record={}), True), + "current, mapped": ( + dict(stamp=year_file_stamp(), record=dict(LEGACY_INPUT_RENAMES)), + True, + ), + "6.1.1 and earlier": (dict(), False), + "record only": (dict(record={}), False), + "missing entry": (dict(stamp=_stamp({}), record={}), False), + "extra entry": ( + dict(stamp=_stamp({**LEGACY_INPUT_RENAMES, "old": "new"}), record={}), + False, + ), + "retargeted": (dict(stamp=_stamp({LEGACY: "elsewhere"}), record={}), False), + "older revision": (dict(stamp=_stamp(LEGACY_INPUT_RENAMES, 0), record={}), False), + "newer revision": (dict(stamp={"revision": 2}, record={}), False), + "malformed stamp": (dict(stamp=("text", "not json"), record={}), False), + "no record": (dict(stamp=year_file_stamp()), False), + "record outside": (dict(stamp=year_file_stamp(), record={"x": "y"}), False), + "not HDF5": (None, False), + "missing": (None, None), +} + + +def _make_state(path: Path, state: str) -> None: + written, _ = STATES[state] + if state == "missing": + return + if state == "not HDF5": + path.parent.mkdir(parents=True, exist_ok=True) + path.write_text("not an H5 file") + return + _write_year_file(path, **written) + + +def _create_current_files(datasets, years, data_folder, allow_unmanaged=False): + """Stand in for ``create_datasets``: write current year files.""" + for year in years: + _write_year_file( + Path(data_folder) / f"{STEM}_year_{year}.h5", + stamp=year_file_stamp(), + record={}, + ) + return {"created": True} + + +@settings(max_examples=60, deadline=None) +@given(states=st.lists(st.sampled_from(sorted(STATES)), min_size=1, max_size=3)) +def test_ensure_cuts_again_exactly_when_load_would_refuse(states): + years = [2024 + index for index in range(len(states))] + stale = [state for state in states if STATES[state][1] is False] + # load_datasets stops at the first year file it cannot use. + first_unusable = next((s for s in states if STATES[s][1] is not True), None) + with tempfile.TemporaryDirectory() as directory: + data_folder = Path(directory) / "data" + for year, state in zip(years, states): + _make_state(data_folder / f"{STEM}_year_{year}.h5", state) + + if first_unusable is not None and STATES[first_unusable][1] is False: + with pytest.raises(ValueError, match="is out of date: "): + load_datasets( + datasets=[SOURCE], years=years, data_folder=str(data_folder) + ) + + expected_warnings = sorted( + (str(path), problem) + for year in years + if (path := data_folder / f"{STEM}_year_{year}.h5").exists() + and (problem := year_file_stamp_problem(path)) is not None + ) + create = MagicMock(side_effect=_create_current_files) + with ( + patch.object(us_datasets, "create_datasets", create), + patch.object(us_datasets.logger, "warning") as warning, + ): + ensure_datasets( + datasets=[SOURCE], years=years, data_folder=str(data_folder) + ) + + assert create.called == (bool(stale) or "missing" in states) + # One warning per out-of-date file, naming it and its reason. + assert warning.call_count == len(stale) + assert ( + sorted( + (str(call.args[1]), call.args[2]) for call in warning.call_args_list + ) + == expected_warnings + ) + for call in warning.call_args_list: + assert call.args[0] == ( + "US year file %s is out of date: %s. Creating the year files again." + ) + + # The year files are now current, so they are reused. + create.reset_mock() + loaded = ensure_datasets( + datasets=[SOURCE], years=years, data_folder=str(data_folder) + ) + create.assert_not_called() + assert sorted(loaded) == [f"{STEM}_{year}" for year in years] + + +@pytest.mark.parametrize( + "states", [["current", "current, mapped"], ["current", "record only"]] +) +def test_ensure_decides_from_the_stamps_alone(tmp_path, states): + """Deciding never opens a year file's tables or the source dataset.""" + data_folder = tmp_path / "data" + years = [2024, 2025] + for year, state in zip(years, states): + _make_state(data_folder / f"{STEM}_year_{year}.h5", state) + load, create = MagicMock(return_value="loaded"), MagicMock(return_value="created") + + def refuse(*args, **kwargs): + raise AssertionError("deciding opened a table or the source dataset") + + class RefusedStore: + """A class, so pandas' ``isinstance(..., HDFStore)`` checks still work.""" + + def __init__(self, *args, **kwargs): + refuse() + + real_file = h5py.File + + class StampAndRecordOnly: + """An H5 file that lets the decision read only the stamp and record.""" + + 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 (YEAR_FILE_STAMP_H5_DATASET, RENAMES_H5_DATASET), name + return self._file.get(name, default) + + with ( + patch.object(h5py, "File", StampAndRecordOnly), + patch.object(pd, "HDFStore", RefusedStore), + # ``pd.read_hdf`` looks HDFStore up in its own module, and PyTables + # opens files without h5py. + patch("pandas.io.pytables.HDFStore", RefusedStore), + patch("tables.open_file", refuse), + patch.object(us_datasets, "materialize_dataset", refuse), + patch.object(us_datasets, "load_datasets", load), + patch.object(us_datasets, "create_datasets", create), + ): + result = ensure_datasets( + datasets=[SOURCE], years=years, data_folder=str(data_folder) + ) + + stale = any(STATES[state][1] is False for state in states) + assert result == ("created" if stale else "loaded") + + +def test_the_refusal_names_the_file_and_what_it_may_lack(tmp_path): + data_folder = tmp_path / "data" + path = data_folder / f"{STEM}_year_2024.h5" + _write_year_file(path, stamp=_stamp({}), record={}) + + with pytest.raises(ValueError) as refused: + load_datasets(datasets=[SOURCE], years=[2024], data_folder=str(data_folder)) + + assert str(refused.value) == ( + f"US year file {path} is out of date: it was cut before these renames " + f"were registered, so it may have lost them: {LEGACY} -> {LIVE}. " + "Regenerate it with ensure_datasets() or create_datasets()." + )