diff --git a/changelog.d/600.fixed.md b/changelog.d/600.fixed.md new file mode 100644 index 00000000..f471a937 --- /dev/null +++ b/changelog.d/600.fixed.md @@ -0,0 +1 @@ +UK year files now record the source dataset they were cut from, the policyengine-core and policyengine-uk versions that cut them, and the year they were cut for, as US year files record their source and versions. `ensure_datasets` creates a UK year file again, and `load_datasets` refuses it, when the requested dataset now resolves to another source, another version is installed, or the file holds another year's tables, so re-pinning a dataset such as `enhanced_frs_2024_25` or `populace_uk_2023` under the same name, editing a local source, or upgrading either package no longer reuses year files cut before. Deciding reads only each year file's stamp and whether it records its data year, and never downloads or reads the source. `create_datasets` writes each UK year file beside its path and moves it into place, stamped, only once the source is checked again, so a year file's tables and stamp always come from one cut; if the source changed, it writes none of that dataset's year files and raises `DatasetMaterializationError`. Requesting two different sources that share a name (and so would share year files) now raises `ValueError`. Year files written before this release are created again once. The stamp code the US and UK share now lives in `policyengine.provenance.year_file_stamp`. diff --git a/data/county_fips_2020.csv.gz b/data/county_fips_2020.csv.gz new file mode 100644 index 00000000..b8cd43c6 Binary files /dev/null and b/data/county_fips_2020.csv.gz differ diff --git a/docs/microsim.md b/docs/microsim.md index aef5afd3..ac5a0500 100644 --- a/docs/microsim.md +++ b/docs/microsim.md @@ -263,12 +263,99 @@ poverty median, and `shareholding`, which spreads corporate taxes across households (PolicyEngine/policyengine.py#567). Variables of a person or benefit unit that do not depend on those match the national run. -Year files written by earlier releases have no recorded data year. -`ensure_datasets` writes them again and `load_datasets` refuses them. A year -file opened directly, as `PolicyEngineUKDataset(filepath=...)`, is not -checked. A saved UK simulation output records the data year its run anchored -on; `Simulation.load()` refuses an output saved without one, and -`Simulation.ensure()` runs it again. +UK year files are named after the dataset's logical name, +`{data_folder}/{name}_year_{year}.h5`, which drops the revision, so files cut +from every release of a dataset share names. `create_datasets` stamps each +with what it was cut from and with, as UTF-8 JSON in +`uk.datasets.YEAR_FILE_STAMP_H5_DATASET` (`policyengine_uk_year_file_stamp`). +US year files store theirs under another name, so a year file of one country +never passes as one of the other. + +```json +{ + "packages": {"policyengine-core": "3.33.0", "policyengine-uk": "2.126.12"}, + "revision": 1, + "source": { + "sha256": "5ba399ea5ab9a6c7d4e9aaf6c623080beb34d9ad5f330f6244dbad72b8d960b6", + "uri": "hf://policyengine/policyengine-uk-data-private/enhanced_frs_2024_25.h5@1.58.0" + }, + "year": 2026 +} +``` + +- `packages`: the installed policyengine-core and policyengine-uk versions. + policyengine-uk extends a single-year source to later years with + `extend_single_year_dataset`, which uprates each year from the one before + by policyengine-uk parameters read through policyengine-core's parameter + tree, so either version can change a projected file's values; +- `revision`: the UK year-file format revision; +- `source`: the identity of the dataset it was cut from, in the three shapes + US year files use (see [Year files](#year-files)): a bundle dataset's + pinned `hf://…@revision` URI and verified sha256, an explicit Hugging Face + URI as given, or a local file's resolved path, size and modification time; +- `year`: the year whose tables the file holds, so a year file copied or + renamed to another year's name is not taken for that year. + +`ensure_datasets` reuses a UK year file only when it records its data year +and its stamp equals the one a file cut now from the requested dataset, for +the requested year, would carry. Otherwise it logs why and creates the year +files again. `load_datasets` refuses the file with the same reason, which +names the year, the source or the package versions it was cut for, from or +with, and the ones requested or installed. +`uk.datasets.year_file_stamp_problem(path, source, year)` returns the +reason, or `None` for a current year file; `source` is +`dataset_materialization.requested_dataset_source_identity("uk", dataset)`. + +So a re-pinned dataset (a new release of `enhanced_frs_2024_25` or +`populace_uk_2023` under the same name), an edited local source, or a new +policyengine-uk or policyengine-core version makes the affected year files +out of date. Year files written by earlier releases have no stamp, and those +written before policyengine.py kept the data year have no data year either: +`ensure_datasets` writes them again once, and `load_datasets` refuses them. + +Deciding reads only each year file's stamp and whether it has a data year, +never its tables. As for US year files, the requested source is identified +without downloading or reading it, and package versions come from the +installed distributions' metadata, so policyengine-uk is not imported. + +`create_datasets` writes each year file to a hidden temporary file beside its +path (`.{name}_year_{year}.…partial`). It records the source file's size and +modification time right after materializing it, and checks them again +(re-hashing a downloaded file) once a dataset's year files are written. Only +then does it stamp each temporary file and move it onto its path, in one +step. The tables and the stamp at a year file's path therefore always come +from one cut, even when two processes cut the same year files at once. If the +source changed, none of that dataset's year files is moved into place, the +ones already there are left as they were, and `DatasetMaterializationError` +is raised. A process killed mid-cut can leave a temporary file behind, which +is never read and can be deleted. Requested datasets that are different +sources with one name are refused with `ValueError`. + +Moving a year file into place replaces the file at its path rather than +rewriting it. The cut therefore needs write access to the data folder (and, +in a folder with the sticky bit set, ownership of the year file it replaces +or of the folder), and the year file gets the mode and group a file the +cutting user creates there would get, not those of the file it replaces. + +The limits listed under [Year files](#year-files) apply to UK year files too, +with policyengine-uk in place of policyengine-us, except that UK year files +are moved into place whole. Three more apply: + +- the versions are those installed when the files are cut. A package upgraded + in a running process after policyengine-uk was imported is recorded at its + new version, though the cut ran the code already imported. Restart the + process after upgrading; +- policyengine-uk draws new student loan holders in each projected year from + numpy's random generator, seeded by the year. numpy's version is not + recorded; +- a multi-year source is read as stored, without the projection. + +A year file opened directly, as `PolicyEngineUKDataset(filepath=...)`, is not +checked, and saving it again with `save()` drops its stamp. A saved UK +simulation output records the data year its run anchored on; +`Simulation.load()` refuses an output saved without one, and +`Simulation.ensure()` runs it again. A saved UK output does not record the +stamp of the year file it was calculated from. ## Simulations diff --git a/src/policyengine/provenance/year_file_stamp.py b/src/policyengine/provenance/year_file_stamp.py new file mode 100644 index 00000000..61d3af3c --- /dev/null +++ b/src/policyengine/provenance/year_file_stamp.py @@ -0,0 +1,275 @@ +"""Stamp year files with what they were cut from and under. + +A country's ``create_datasets`` cuts one file per requested year from a +source dataset and names it ``{data_folder}/{stem}_year_{year}.h5``, after +the source's logical name. The name drops any revision, so files cut from +every release of a dataset share names. Each year file is therefore stamped +with the identity of the source it was cut from and the versions of the +packages that calculated it, and the country's ``ensure_datasets`` and +``load_datasets`` compare that stamp with the one a file cut now would carry. + +Each country stores its stamp, as UTF-8 JSON, in an H5 dataset of its own +name, so a year file of one country never passes as one of another: +``tax_benefit_models.us.legacy_inputs.YEAR_FILE_STAMP_H5_DATASET`` (which +also records the register of renamed stored inputs) and +``tax_benefit_models.uk.datasets.YEAR_FILE_STAMP_H5_DATASET`` (which also +records the year the file was cut for). + +This module holds what the countries share: + +- reading and writing a JSON value in an H5 dataset + (:func:`read_json_dataset`, :func:`read_stored_json`, + :func:`write_json_dataset`); +- the installed versions of the packages a stamp records + (:func:`installed_package_versions`); +- the clauses naming how a stamp's source and package versions differ from + the current ones (:func:`source_and_package_differences`); +- each requested dataset's year-file stem and source identity, refusing two + sources that would share year files (:func:`year_file_sources`); +- writing a year file beside its path and moving it into place only once it + is stamped (:func:`temporary_year_file`, :func:`publish_year_file`). +""" + +from __future__ import annotations + +import json +import os +import secrets +from collections.abc import Iterable, Mapping +from importlib import metadata as importlib_metadata +from pathlib import Path +from typing import Any, Optional + +from .dataset_materialization import ( + describe_dataset_source_identity, + requested_dataset_source_identity, +) +from .manifest import dataset_logical_name, resolve_dataset_reference + +#: 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() + +#: Errors opening or reading a year file that say nothing about its +#: contents: no permission, a directory at its path, or a lock another +#: process holds. A year-file check raises them rather than reporting the +#: file out of date, since the file could not be written again either. +UNRELATED_TO_CONTENTS = (PermissionError, IsADirectoryError, BlockingIOError) + +#: The clause a year-file check gives for a file h5py cannot open. +NOT_HDF5 = "it cannot be read as an HDF5 file (it may be half-written or corrupt)" + + +def installed_package_versions(names: Iterable[str]) -> dict[str, Optional[str]]: + """Return the installed version of each named distribution. + + ``None`` stands for a distribution that is not installed. The versions + come from the installed distributions' metadata, so no package is + imported. + """ + versions: dict[str, Optional[str]] = {} + for name in names: + try: + versions[name] = importlib_metadata.version(name) + except importlib_metadata.PackageNotFoundError: + versions[name] = None + return versions + + +def read_json_dataset(file: Any, name: str) -> Any: + """Return the JSON value stored in an open H5 file's ``name`` dataset. + + Returns :data:`ABSENT` when the file has no such node, and + :data:`UNREADABLE` when the node is not a scalar dataset holding UTF-8 + JSON, or cannot be followed. Errors in :data:`UNRELATED_TO_CONTENTS` are + raised. + """ + import h5py + + try: + # A cycle of soft links, or a chain too long to follow, raises + # RuntimeError here. + 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, a RuntimeError. + return json.loads(node.asstr()[()]) + except UNRELATED_TO_CONTENTS: + raise + except (KeyError, OSError, RuntimeError, 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 name in file: + del file[name] + # UTF-8 JSON in a dataset, as for the SPM receipt of a saved output. + file.create_dataset( + name, + data=json.dumps(value, sort_keys=True), + dtype=h5py.string_dtype("utf-8"), + ) + + +def read_stored_json(path: str | Path, name: str) -> Any | None: + """Return the JSON value of an H5 file's ``name`` dataset, well formed or + not. + + ``None`` means the file stores no readable JSON there (see + :func:`read_json_dataset`) or is not an H5 file. + """ + import h5py + + try: + with h5py.File(path, "r") as file: + value = read_json_dataset(file, name) + except OSError: + return None + return None if value is ABSENT or value is UNREADABLE else value + + +def temporary_year_file(path: str | Path) -> Path: + """Create an empty file beside ``path`` to write a year file to. + + ``create_datasets`` writes a year file's tables there, stamps it once the + source has been checked again, and only then moves it onto ``path`` (see + :func:`publish_year_file`). The tables and the stamp at ``path`` then + always come from one cut, even when two processes cut the same year files + at once, and a cut that stops part-way leaves ``path`` as it was. The + file's name starts with ``.{stem}.`` and ends with ``.partial``; the + caller removes it if the cut fails. + + The file is created with the mode a new file gets (``0o666`` less the + umask), as pandas would create the year file itself, so a year file moved + into place is as readable by others as one written at its path. (A file + from ``tempfile.mkstemp`` would be readable by its owner alone.) + """ + path = Path(path) + path.parent.mkdir(parents=True, exist_ok=True) + flags = os.O_WRONLY | os.O_CREAT | os.O_EXCL + for _ in range(100): + temporary = path.parent / f".{path.stem}.{secrets.token_hex(8)}.partial" + try: + descriptor = os.open(temporary, flags, 0o666) + except FileExistsError: + continue + os.close(descriptor) + return temporary + raise FileExistsError(f"Could not create a temporary file beside {path}") + + +def publish_year_file( + temporary: str | Path, path: str | Path, name: str, stamp: Mapping[str, Any] +) -> None: + """Stamp a year file written to ``temporary``, then move it onto ``path``. + + ``name`` is the country's stamp dataset. The move replaces any file at + ``path`` in one step (``os.replace`` within one folder), so a reader sees + either the old file or the new one, each with its own stamp. + """ + write_json_dataset(temporary, name, dict(stamp)) + os.replace(temporary, path) + + +def is_version_map(value: Any) -> bool: + """Whether ``value`` maps package names to versions, ``None`` for absent.""" + return isinstance(value, dict) and all( + isinstance(key, str) and (item is None or isinstance(item, str)) + for key, item in value.items() + ) + + +def source_and_package_differences( + stamp: Mapping[str, Any], current: Mapping[str, Any] +) -> list[str]: + """Return the clauses naming how a stamp's source and packages differ. + + ``stamp`` is a year file's well-formed stamp and ``current`` the stamp a + file cut now would carry: each has a ``"source"`` that + ``dataset_materialization.is_dataset_source_identity`` accepts and + ``"packages"`` that :func:`is_version_map` accepts. One clause, starting + "it", names both sources when they differ, and one more names each + package whose version differs, in name order. A package the stamp + records that ``current`` does not is named too, and a package recorded + in both with one version is not. + + Returns no clause exactly when the sources and the package maps are + equal. + """ + reasons = [] + if stamp["source"] != current["source"]: + reasons.append( + "it was cut from " + f"{describe_dataset_source_identity(stamp['source'])}, and the " + "requested source is " + f"{describe_dataset_source_identity(current['source'])}" + ) + stamped_packages, packages = stamp["packages"], current["packages"] + for name in sorted(stamped_packages.keys() | packages.keys()): + if name not in packages: + reasons.append( + f"it records the {name} version it was cut with, which this " + "version of policyengine.py does not check" + ) + continue + if name in stamped_packages and stamped_packages[name] == packages[name]: + continue + cut_with = ( + f"{name} {stamped_packages[name]}" + if stamped_packages.get(name) is not None + else f"no recorded {name} version" + ) + installed = ( + f"{name} {packages[name]} is installed" + if packages[name] is not None + else f"{name} is not installed" + ) + reasons.append(f"it was cut with {cut_with}, and {installed}") + return reasons + + +def year_file_sources( + country_id: str, datasets: list[str], allow_unmanaged: bool = True +) -> list[tuple[str, dict[str, Any]]]: + """Return each requested dataset's year-file stem and source identity. + + The stem names the year files (``{stem}_year_{year}.h5``) and drops any + revision, so files cut from every release of a logical dataset share + names; the identity tells them apart. Neither downloads nor reads the + source (see ``dataset_materialization.requested_dataset_source_identity``). + A request ``materialize_dataset`` would refuse with the same + ``allow_unmanaged`` raises the same error, before its stem is resolved. + + Raises ``ValueError`` when two of them are different sources with one + stem. They would write the same year files, each cut overwriting the + other's, so a request of both would cut them all again every time. + """ + sources = [] + for dataset in datasets: + identity = requested_dataset_source_identity( + country_id, dataset, allow_unmanaged=allow_unmanaged + ) + stem = dataset_logical_name(resolve_dataset_reference(country_id, dataset)) + sources.append((stem, identity)) + first_by_stem: dict[str, tuple[str, dict[str, Any]]] = {} + for dataset, (stem, identity) in zip(datasets, sources): + first, first_identity = first_by_stem.setdefault(stem, (dataset, identity)) + if first_identity != identity: + raise ValueError( + f"Datasets {first!r} and {dataset!r} would share the " + f"{country_id.upper()} year files {stem}_year_{{year}}.h5, but " + "are different sources: " + f"{describe_dataset_source_identity(first_identity)} and " + f"{describe_dataset_source_identity(identity)}. Request them " + "in separate calls with separate data folders." + ) + return sources diff --git a/src/policyengine/tax_benefit_models/uk/datasets.py b/src/policyengine/tax_benefit_models/uk/datasets.py index dbe31dbc..4efe38e3 100644 --- a/src/policyengine/tax_benefit_models/uk/datasets.py +++ b/src/policyengine/tax_benefit_models/uk/datasets.py @@ -1,29 +1,75 @@ +import logging +from collections.abc import Mapping from pathlib import Path -from typing import Optional +from typing import Any, Optional import pandas as pd from microdf import MicroDataFrame from pydantic import ConfigDict from policyengine.core import Dataset, YearData +from policyengine.provenance import year_file_stamp as _shared_stamp from policyengine.provenance.dataset_materialization import ( + DatasetMaterializationError, + dataset_source_change, + dataset_source_identity, + dataset_source_snapshot, + is_dataset_source_identity, materialize_dataset, ) -from policyengine.provenance.manifest import ( - dataset_logical_name, - get_release_manifest, - resolve_dataset_reference, +from policyengine.provenance.manifest import get_release_manifest +from policyengine.provenance.year_file_stamp import ( + ABSENT, + NOT_HDF5, + UNRELATED_TO_CONTENTS, + is_version_map, + publish_year_file, + read_json_dataset, + read_stored_json, + source_and_package_differences, + temporary_year_file, + write_json_dataset, + year_file_sources, ) +logger = logging.getLogger(__name__) + #: HDF5 key holding the data year of a UK year file: the year of the observed #: data the file's tables were projected from (see #: ``PolicyEngineUKDataset.data_year``). ``create_datasets`` writes it into #: every year file. Year files written before policyengine.py kept the data #: year have none: policyengine-uk would take their projected tables as #: observed data, so ``ensure_datasets`` regenerates them and -#: ``load_datasets`` refuses them. +#: ``load_datasets`` refuses them (see :func:`year_file_stamp_problem`). DATA_YEAR_KEY = "data_year" +#: H5 dataset stamping a UK ``create_datasets`` year file, as UTF-8 JSON (see +#: :func:`year_file_stamp`). Its name differs from the US year-file stamp's +#: (``us.legacy_inputs.YEAR_FILE_STAMP_H5_DATASET``), so a year file of one +#: country never passes as one of the other. ``PolicyEngineUKDataset.save()`` +#: neither writes nor keeps it, so saving a year file again drops the stamp. +YEAR_FILE_STAMP_H5_DATASET = "policyengine_uk_year_file_stamp" + +#: Format revision of UK ``create_datasets`` year files, stored in their +#: stamp (see :func:`year_file_stamp`). Raise it when the stamp changes +#: shape, or when a change to year-file extraction means that files cut +#: before it must be cut again, for a reason the rest of the stamp does not +#: capture. +#: +#: Revision 1 is the first: UK year files cut before it have no stamp, so +#: they are cut again once. +YEAR_FILE_STAMP_REVISION = 1 + +#: Distributions whose installed versions a UK year file's stamp records. +#: ``create_datasets`` stores each year's tables as policyengine-uk projects +#: the source dataset. policyengine-uk extends a single-year source to later +#: years with ``extend_single_year_dataset``, which uprates each year from the +#: one before by policyengine-uk parameters read through policyengine-core's +#: parameter tree, so either can change a year file's values. (It also draws +#: new student loan holders in each projected year from numpy's random +#: generator, seeded by the year; numpy's version is not recorded.) +YEAR_FILE_PACKAGES = ("policyengine-core", "policyengine-uk") + #: Prefix of the HDF5 keys holding the data year's entity tables in a year #: file whose year comes after its data year. DATA_YEAR_TABLE_PREFIX = "data_year_" @@ -140,11 +186,13 @@ def save(self) -> None: if not filepath.parent.exists(): filepath.parent.mkdir(parents=True, exist_ok=True) + # Mode "w" truncates the file first, so saving a year file again + # drops the stamp ``create_datasets`` wrote after it. with pd.HDFStore(filepath, mode="w") as store: _put_year_data(store, self.data) if self.data_year_data is not None: _put_year_data(store, self.data_year_data, DATA_YEAR_TABLE_PREFIX) - # Written last: its presence marks a complete file. + # Written last: its presence marks complete tables. if self.data_year is not None: store.put(DATA_YEAR_KEY, pd.Series([int(self.data_year)])) @@ -263,12 +311,252 @@ def _year_file_records_data_year(path: Path) -> bool: return f"/{DATA_YEAR_KEY}" in store.keys() +def installed_package_versions() -> dict[str, Optional[str]]: + """Return the installed version of each of :data:`YEAR_FILE_PACKAGES`. + + ``None`` stands for a package that is not installed. The versions come + from the installed distributions' metadata, so no package is imported. + """ + return _shared_stamp.installed_package_versions(YEAR_FILE_PACKAGES) + + +def year_file_stamp( + source: Mapping[str, Any], + year: int, + packages: Optional[Mapping[str, Optional[str]]] = None, +) -> dict[str, Any]: + """Return the stamp for a UK year file for ``year`` cut now from ``source``. + + ``create_datasets`` writes it to each year file, after the tables and the + data year, as UTF-8 JSON in :data:`YEAR_FILE_STAMP_H5_DATASET`:: + + {"packages": {"policyengine-core": str | null, + "policyengine-uk": str | null}, + "revision": int, + "source": {...}, + "year": int} + + ``source`` is the source dataset's identity (see + ``policyengine.provenance.dataset_materialization.dataset_source_identity``), + and ``packages`` defaults to :func:`installed_package_versions`, read + when this is called (``null`` for a package that is not installed). + ``revision`` is :data:`YEAR_FILE_STAMP_REVISION`, and ``year`` the year + whose tables the file holds, so a year file copied or renamed to another + year's name is not taken for that year. + """ + if packages is None: + packages = installed_package_versions() + return { + "packages": dict(sorted(packages.items())), + "revision": YEAR_FILE_STAMP_REVISION, + "source": dict(source), + "year": int(year), + } + + +def read_year_file_stamp(path: str | Path) -> Any | None: + """Return the JSON value of a UK 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. + """ + return read_stored_json(path, YEAR_FILE_STAMP_H5_DATASET) + + +def write_year_file_stamp(path: str | Path, stamp: Mapping[str, Any]) -> None: + """Stamp a UK year file (see :func:`year_file_stamp`), replacing any stamp.""" + write_json_dataset(path, YEAR_FILE_STAMP_H5_DATASET, dict(stamp)) + + +#: What a run of a year file without a data year would do wrong. +_WITHOUT_A_DATA_YEAR = ( + "policyengine-uk would take its projected tables as observed data and " + "uprate the State Pension by CPI rather than the triple lock" +) + +#: Why an unstamped UK year file without a recorded data year is out of date. +NO_DATA_YEAR = ( + "it has no recorded data year, so it was written before policyengine.py " + "kept the observed year its tables were projected from, or not written " + f"completely; {_WITHOUT_A_DATA_YEAR}" +) + +#: Why a stamped UK year file without a recorded data year is out of date. +STAMPED_WITHOUT_DATA_YEAR = ( + "it is stamped but records no data year, which every year file " + f"create_datasets writes records; {_WITHOUT_A_DATA_YEAR}" +) + + +def year_file_stamp_problem( + path: str | Path, + source: Mapping[str, Any], + year: int, + packages: Optional[Mapping[str, Optional[str]]] = None, +) -> Optional[str]: + """Return why a UK ``create_datasets`` year file is out of date, or ``None``. + + ``source`` is the identity of the dataset the year file is requested for + (see ``policyengine.provenance.dataset_materialization + .requested_dataset_source_identity``), ``year`` the year it is requested + for, and ``packages`` defaults to :func:`installed_package_versions`. A + year file is current only when both of these hold: + + - it records its data year (``DATA_YEAR_KEY``). Without one, + policyengine-uk would take its projected tables as observed data; + - its stamp equals ``year_file_stamp(source, year, packages)``: the same + format revision, year, source identity and package versions. A file cut + from another revision or file of a dataset that keeps its logical name + (and so the year file's name) holds the other data, and the package + versions projected every table in the file. + + Only the stamp and the presence of the data year are read, never the + tables or the source dataset. A file that cannot be read as HDF5 is out of + date, and so is one whose stamp or data year is a link that cannot be + followed. 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. It names the source + it was cut from and the source requested, and the package versions it + was cut with and those installed, wherever they differ. + """ + import h5py + + current = year_file_stamp(source, year, packages) + try: + with h5py.File(path, "r") as file: + stamp = read_json_dataset(file, YEAR_FILE_STAMP_H5_DATASET) + records_data_year = _records_data_year(file) + except UNRELATED_TO_CONTENTS: + raise + except OSError: + return NOT_HDF5 + + # The revision comes first: a later revision may shape its stamp, or + # record the data year, 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 not records_data_year: + return NO_DATA_YEAR if stamp is ABSENT else STAMPED_WITHOUT_DATA_YEAR + if stamp is ABSENT: + return ( + "it has no stamp of the source dataset and package versions it was " + "cut from and with, so it was cut by an older policyengine.py or " + "not written completely" + ) + # ``type(...) is int`` rather than ``isinstance``: JSON ``true`` loads as + # ``True``, which equals 1. + if not ( + isinstance(stamp, dict) + and set(stamp) == {"packages", "revision", "source", "year"} + and type(revision) is int + and type(stamp["year"]) is int + and is_version_map(stamp["packages"]) + and is_dataset_source_identity(stamp["source"]) + ): + return "its stamp of what it was cut from and with is malformed" + if stamp == current: + return None + reasons = [] + if stamp["year"] != current["year"]: + reasons.append(f"it was cut for {stamp['year']}, not {current['year']}") + reasons.extend(source_and_package_differences(stamp, current)) + return "; ".join(reasons) + + +def _records_data_year(file: Any) -> bool: + """Return whether an open H5 year file holds its data year. + + pandas stores the data year (``DATA_YEAR_KEY``) as a group. A node that + is absent, not a group, or a link that cannot be followed is no data + year. Errors in ``UNRELATED_TO_CONTENTS`` are raised. + """ + import h5py + + try: + # A cycle of soft links, or a chain too long to follow, raises + # RuntimeError here. + return isinstance(file.get(DATA_YEAR_KEY), h5py.Group) + except UNRELATED_TO_CONTENTS: + raise + except (KeyError, OSError, RuntimeError): + return False + + +def _year_file_sources( + datasets: list[str], allow_unmanaged: bool = True +) -> list[tuple[str, dict[str, Any]]]: + """Return each requested dataset's year-file stem and source identity. + + Neither downloads nor reads the source. Raises ``ValueError`` when two of + them are different sources with one stem, since they would overwrite each + other's year files (see ``year_file_stamp.year_file_sources``). + """ + return year_file_sources("uk", datasets, allow_unmanaged) + + def create_datasets( datasets: Optional[list[str]] = None, years: list[int] = [2026, 2027, 2028, 2029, 2030], data_folder: str = "./data", allow_unmanaged: bool = False, ) -> dict[str, PolicyEngineUKDataset]: + """Cut UK year files from logical dataset names, URIs or local files. + + Each year file is written to ``{data_folder}/{dataset_stem}_year_{year}.h5`` + with its tables, its data year and, for a year after the data year, the + data year's tables. Each is stamped with the identity of the source + dataset it was cut from, the installed policyengine-core and + policyengine-uk versions, its year and the year-file format revision (see + :func:`year_file_stamp`), which ``ensure_datasets`` and ``load_datasets`` + check. + + Each year file is written to a temporary file beside its path (see + ``year_file_stamp.temporary_year_file``). Once a dataset's year files are + all written, the source is checked again, and only then is each stamped + and moved onto its path. The tables and the stamp at a year file's path + therefore always come from one cut, even when two processes cut the same + year files at once, and a cut that fails leaves the year files already + there as they were. + + The source file's size and modification time are recorded right after it + is materialized, and checked again once the year files are written. A + downloaded file is also hashed again, against the sha256 a bundle + dataset was verified against or an explicit URI was downloaded with (see + ``dataset_materialization.dataset_source_change``). Another process can + replace a file in a shared data folder, for example by materializing + another release of the same dataset there. If the source changed, none of + that dataset's year files is moved into place (those of datasets cut + before it in the same call already are) and + ``DatasetMaterializationError`` is raised. + Two requested datasets that are different sources with one stem are + refused with ``ValueError`` before anything is cut, since they would + overwrite each other's year files. + + Args: + datasets: Logical dataset names, Hugging Face URIs or local paths. + ``None`` selects the certified default dataset. + years: Years to cut a year file for. + data_folder: Directory to write the year files (and keep downloaded + sources) in. + allow_unmanaged: Allow explicit URIs and local paths that bypass the + release bundle. + + Returns: + Dictionary mapping dataset keys (e.g. ``"enhanced_frs_2024_25_2026"``) + to PolicyEngineUKDataset objects. + """ + if datasets is not None and len(datasets) > 1: + _year_file_sources(datasets, allow_unmanaged) + dataset_requests: list[Optional[str]] = [None] if datasets is None else datasets result = {} for dataset in dataset_requests: @@ -279,6 +567,12 @@ def create_datasets( data_dir=Path(data_folder), ) dataset_stem = source.name + # What each year file is cut from and under: the source, identified + # before it is loaded, and the package versions. + identity = dataset_source_identity(source) + # The source is checked again only when year files are cut from it. + snapshot = dataset_source_snapshot(source) if years else None + packages = installed_package_versions() from policyengine_uk import Microsimulation sim = Microsimulation(dataset=source.path) @@ -290,21 +584,56 @@ def create_datasets( # file anchors on the same observed year as a run of the source. data_year = int(min(sim.dataset.years)) data_year_data = _year_data_with_weights(sim.dataset[data_year]) - for year in years: - uk_dataset = PolicyEngineUKDataset( - id=f"{dataset_stem}_year_{year}", - name=f"{dataset_stem}-year-{year}", - description=f"UK Dataset for year {year} based on {dataset_stem}", - filepath=f"{data_folder}/{dataset_stem}_year_{year}.h5", - year=int(year), - data=_year_data_with_weights(sim.dataset[year]), - data_year=data_year, - data_year_data=data_year_data if int(year) > data_year else None, + # (temporary file, year file, year, result key) for each year file + # written. + cut: list[tuple[Path, str, int, str]] = [] + try: + for year in years: + filepath = f"{data_folder}/{dataset_stem}_year_{year}.h5" + dataset_key = f"{dataset_stem}_{year}" + temporary = temporary_year_file(filepath) + cut.append((temporary, filepath, int(year), dataset_key)) + uk_dataset = PolicyEngineUKDataset( + id=f"{dataset_stem}_year_{year}", + name=f"{dataset_stem}-year-{year}", + description=f"UK Dataset for year {year} based on {dataset_stem}", + filepath=str(temporary), + year=int(year), + data=_year_data_with_weights(sim.dataset[year]), + data_year=data_year, + data_year_data=data_year_data if int(year) > data_year else None, + ) + uk_dataset.save() + result[dataset_key] = uk_dataset + + # Move the year files into place only if the source file is still + # the one materialized (and, if it was downloaded, still hashes to + # the sha256 it was verified against or downloaded with), so no + # stamp names a source the files may not have been cut from. + change = ( + dataset_source_change(source, identity, snapshot) + if snapshot is not None + else None ) - uk_dataset.save() - - dataset_key = f"{dataset_stem}_{year}" - result[dataset_key] = uk_dataset + if change is not None: + raise DatasetMaterializationError( + f"The source of the UK year files {data_folder}/" + f"{dataset_stem}_year_{{year}}.h5 changed while they were " + f"cut: {change}. None of them was written, so " + "ensure_datasets() cuts them again." + ) + # Stamp last, then move each year file onto its path in one step. + for temporary, filepath, year, dataset_key in cut: + publish_year_file( + temporary, + filepath, + YEAR_FILE_STAMP_H5_DATASET, + year_file_stamp(identity, year, packages), + ) + result[dataset_key].filepath = filepath + finally: + for temporary, *_ in cut: + temporary.unlink(missing_ok=True) return result @@ -316,31 +645,34 @@ def load_datasets( ) -> dict[str, PolicyEngineUKDataset]: """Load UK year files written by ``create_datasets``. - A year file without a recorded data year (``DATA_YEAR_KEY``) was written - before policyengine.py kept the observed year its tables were projected - from. policyengine-uk would take those projected tables as observed data - and uprate the State Pension by CPI rather than the triple lock, so such - a file is refused. + A year file is refused unless it records its data year and + ``create_datasets`` stamped it as cut from the source the dataset + resolves to now, with the installed policyengine-core and policyengine-uk + versions and this release's year-file format revision (see + :func:`year_file_stamp_problem`). A year file without a recorded data + year was written before policyengine.py kept the observed year its tables + were projected from: policyengine-uk would take those projected tables as + observed data and uprate the State Pension by CPI rather than the triple + lock. A file cut from an earlier release of a dataset the bundle has since + re-pinned holds that release's data. The error says why, naming the two + sources or the two package versions. Two requested datasets that are + different sources with one stem are refused, since they would share year + files. """ if datasets is None: datasets = [get_release_manifest("uk").default_dataset] result = {} - for dataset in datasets: - resolved_dataset = resolve_dataset_reference("uk", dataset) - dataset_stem = dataset_logical_name(resolved_dataset) + for dataset_stem, source in _year_file_sources(datasets): for year in years: filepath = f"{data_folder}/{dataset_stem}_year_{year}.h5" - if Path(filepath).exists() and not _year_file_records_data_year( - Path(filepath) - ): - raise ValueError( - f"UK year file {filepath} has no recorded data year, so it " - "was written before policyengine.py kept the observed year " - "its tables were projected from. policyengine-uk would " - "take the projected tables as observed data and uprate the " - "State Pension by CPI rather than the triple lock. " - "Regenerate it with ensure_datasets() or create_datasets()." - ) + if Path(filepath).exists(): + problem = year_file_stamp_problem(filepath, source, int(year)) + if problem is not None: + raise ValueError( + f"UK year file {filepath} is out of date: {problem}. " + "Regenerate it with ensure_datasets() or " + "create_datasets()." + ) # Constructing with a filepath loads the file. uk_dataset = PolicyEngineUKDataset( name=f"{dataset_stem}-year-{year}", @@ -363,14 +695,29 @@ def ensure_datasets( ) -> dict[str, PolicyEngineUKDataset]: """Ensure datasets exist, loading if available or creating if not. - Year files without a recorded data year were written before - policyengine.py kept the observed year their tables were projected from - (see ``load_datasets``), so they are created again rather than loaded. + Existing year files are reused only when they record their data year and + ``create_datasets`` stamped them as cut from the source each dataset + resolves to now, with the installed policyengine-core and policyengine-uk + versions and this release's year-file format revision (see + :func:`year_file_stamp_problem`). Only each file's stamp and the presence + of its data year are read to decide, never its tables. The source is + never downloaded or read: a bundle dataset is identified by its manifest + entry, an explicit Hugging Face URI by the URI and a local file by its + file-system metadata (see + ``dataset_materialization.requested_dataset_source_identity``). 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. + Two requested datasets that are different sources with one stem are + refused with ``ValueError``, since they would overwrite each other's year + files. Args: - datasets: List of HuggingFace dataset paths + datasets: List of logical dataset names, Hugging Face URIs or local + paths years: List of years to load/create data for data_folder: Directory containing or to save the dataset files + allow_unmanaged: Allow explicit URIs and local paths that bypass the + release bundle, when the year files are created Returns: Dictionary mapping dataset keys to PolicyEngineUKDataset objects @@ -378,22 +725,27 @@ def ensure_datasets( if datasets is None: datasets = [get_release_manifest("uk").default_dataset] - # Check if all dataset files exist - all_exist = True - for dataset in datasets: - resolved_dataset = resolve_dataset_reference("uk", dataset) - dataset_stem = dataset_logical_name(resolved_dataset) + # Check that every year file exists and is current. A year file without + # a data year, or cut from another source or with other package + # versions, is created again rather than reused. + all_current = True + for dataset_stem, source in _year_file_sources(datasets): for year in years: filepath = Path(f"{data_folder}/{dataset_stem}_year_{year}.h5") - # A year file without a recorded data year would run as observed - # data, so it is regenerated rather than reused. - if not filepath.exists() or not _year_file_records_data_year(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, source, int(year)) + if problem is not None: + all_current = False + logger.warning( + "UK 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/datasets.py b/src/policyengine/tax_benefit_models/us/datasets.py index fd86ecfe..c7ce0cde 100644 --- a/src/policyengine/tax_benefit_models/us/datasets.py +++ b/src/policyengine/tax_benefit_models/us/datasets.py @@ -19,15 +19,12 @@ dataset_source_change, dataset_source_identity, dataset_source_snapshot, - describe_dataset_source_identity, materialize_dataset, - requested_dataset_source_identity, ) from policyengine.provenance.manifest import ( - dataset_logical_name, get_release_manifest, - resolve_dataset_reference, ) +from policyengine.provenance.year_file_stamp import year_file_sources from policyengine.tax_benefit_models.common.model_version import ( build_runtime_dataset_provenance, ) @@ -375,47 +372,16 @@ def _load_policyengine_core_h5(path: Path, year: int) -> USYearData: ) -def _year_file_source( - dataset: str, allow_unmanaged: bool = True -) -> tuple[str, dict[str, Any]]: - """Return a requested dataset's year-file stem and its source identity. - - The stem names the year files (``{stem}_year_{year}.h5``) and drops any - revision, so files cut from every release of a logical dataset share - names; the identity tells them apart. Neither downloads nor reads the - source (see ``requested_dataset_source_identity``). A request - ``materialize_dataset`` would refuse with the same ``allow_unmanaged`` - raises the same error, before the stem is resolved. - """ - identity = requested_dataset_source_identity( - "us", dataset, allow_unmanaged=allow_unmanaged - ) - stem = dataset_logical_name(resolve_dataset_reference("us", dataset)) - return stem, identity - - def _year_file_sources( datasets: list[str], allow_unmanaged: bool = True ) -> list[tuple[str, dict[str, Any]]]: """Return each requested dataset's year-file stem and source identity. - Raises ``ValueError`` when two of them are different sources with one - stem. They would write the same year files, each cut overwriting the - other's, so a request of both would cut them all again every time. + Neither downloads nor reads the source. Raises ``ValueError`` when two of + them are different sources with one stem, since they would overwrite each + other's year files (see ``year_file_stamp.year_file_sources``). """ - sources = [_year_file_source(dataset, allow_unmanaged) for dataset in datasets] - first_by_stem: dict[str, tuple[str, dict[str, Any]]] = {} - for dataset, (stem, identity) in zip(datasets, sources): - first, first_identity = first_by_stem.setdefault(stem, (dataset, identity)) - if first_identity != identity: - raise ValueError( - f"Datasets {first!r} and {dataset!r} would share the US year " - f"files {stem}_year_{{year}}.h5, but are different sources: " - f"{describe_dataset_source_identity(first_identity)} and " - f"{describe_dataset_source_identity(identity)}. Request them " - "in separate calls with separate data folders." - ) - return sources + return year_file_sources("us", datasets, allow_unmanaged) def 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 0b5221bc..c03f78fa 100644 --- a/src/policyengine/tax_benefit_models/us/legacy_inputs.py +++ b/src/policyengine/tax_benefit_models/us/legacy_inputs.py @@ -23,19 +23,37 @@ from __future__ import annotations import inspect -import json from collections.abc import Iterable, Iterator, Mapping -from importlib import metadata as importlib_metadata from pathlib import Path from typing import Any import numpy as np import pandas as pd +from policyengine.provenance import year_file_stamp as _shared_stamp from policyengine.provenance.dataset_materialization import ( - describe_dataset_source_identity, is_dataset_source_identity, ) +from policyengine.provenance.year_file_stamp import ABSENT as _ABSENT +from policyengine.provenance.year_file_stamp import NOT_HDF5 as _NOT_HDF5 +from policyengine.provenance.year_file_stamp import ( + UNRELATED_TO_CONTENTS as _UNRELATED_TO_CONTENTS, +) +from policyengine.provenance.year_file_stamp import ( + is_version_map as _is_version_map, +) +from policyengine.provenance.year_file_stamp import ( + read_json_dataset as _read_json_dataset, +) +from policyengine.provenance.year_file_stamp import ( + read_stored_json as _read_stored_json, +) +from policyengine.provenance.year_file_stamp import ( + source_and_package_differences as _source_and_package_differences, +) +from policyengine.provenance.year_file_stamp import ( + write_json_dataset as _write_json_dataset, +) #: Stored input columns that policyengine-us has since renamed, mapped to the #: live input that replaced each one. @@ -100,6 +118,10 @@ #: now from the requested dataset would carry (see #: :func:`year_file_stamp_problem`). ``PolicyEngineUSDataset.save()`` neither #: writes nor keeps it, so saving a year file again drops the stamp. +#: +#: UK year files store their stamp under a name of their own, so a year file +#: of one country never passes as one of the other (see +#: ``policyengine.provenance.year_file_stamp``). YEAR_FILE_STAMP_H5_DATASET = "policyengine_year_file_stamp" #: Format revision of ``create_datasets`` year files, stored in their stamp. @@ -149,12 +171,6 @@ 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``. @@ -182,13 +198,7 @@ def installed_package_versions() -> dict[str, str | None]: ``None`` stands for a package that is not installed. The versions come from the installed distributions' metadata, so no package is imported. """ - versions: dict[str, str | None] = {} - for name in YEAR_FILE_PACKAGES: - try: - versions[name] = importlib_metadata.version(name) - except importlib_metadata.PackageNotFoundError: - versions[name] = None - return versions + return _shared_stamp.installed_package_versions(YEAR_FILE_PACKAGES) def year_file_stamp( @@ -221,14 +231,7 @@ def read_year_file_stamp(path: str | Path) -> Any | None: ``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 + return _read_stored_json(path, YEAR_FILE_STAMP_H5_DATASET) def write_year_file_stamp(path: str | Path, stamp: Mapping[str, Any]) -> None: @@ -287,10 +290,10 @@ def year_file_stamp_problem( 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): + except _UNRELATED_TO_CONTENTS: raise except OSError: - return "it cannot be read as an HDF5 file (it may be half-written or corrupt)" + return _NOT_HDF5 # The revision comes first: a later revision may shape its stamp and # record differently. @@ -360,88 +363,16 @@ def year_file_stamp_problem( "it was cut under renames this version of policyengine.py does not " f"register: {_describe(unknown)}" ) - if stamp["source"] != current["source"]: - reasons.append( - "it was cut from " - f"{describe_dataset_source_identity(stamp['source'])}, and the " - "requested source is " - f"{describe_dataset_source_identity(current['source'])}" - ) - stamped_packages, packages = stamp["packages"], current["packages"] - for name in sorted(stamped_packages.keys() | packages.keys()): - if name not in packages: - reasons.append( - f"it records the {name} version it was cut with, which this " - "version of policyengine.py does not check" - ) - continue - if name in stamped_packages and stamped_packages[name] == packages[name]: - continue - cut_with = ( - f"{name} {stamped_packages[name]}" - if stamped_packages.get(name) is not None - else f"no recorded {name} version" - ) - installed = ( - f"{name} {packages[name]} is installed" - if packages[name] is not None - else f"{name} is not installed" - ) - reasons.append(f"it was cut with {cut_with}, and {installed}") + reasons.extend(_source_and_package_differences(stamp, current)) 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: - # A cycle of soft links, or a chain too long to follow, raises - # RuntimeError here. - 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, a RuntimeError. - return json.loads(node.asstr()[()]) - except (PermissionError, IsADirectoryError, BlockingIOError): - raise - except (KeyError, OSError, RuntimeError, 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 name in file: - del file[name] - # UTF-8 JSON in a dataset, as for the SPM receipt of a saved output. - file.create_dataset( - 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 _is_version_map(value: Any) -> bool: - """Whether ``value`` maps package names to versions, ``None`` for absent.""" - return isinstance(value, dict) and all( - isinstance(key, str) and (item is None or 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())) diff --git a/tests/test_dataset_runtime.py b/tests/test_dataset_runtime.py index 76cab79b..0dbc546c 100644 --- a/tests/test_dataset_runtime.py +++ b/tests/test_dataset_runtime.py @@ -141,17 +141,12 @@ def test_uk_load_datasets_defaults_to_certified_bundle_dataset(monkeypatch): "_test_policyengine_uk_default_load_datasets", REPO_ROOT / "src/policyengine/tax_benefit_models/uk/datasets.py", ) - resolve = Mock( - return_value=( - "hf://policyengine/policyengine-uk-data-private/" - "enhanced_frs_2024_25.h5@1.56.16" - ) - ) - monkeypatch.setattr(uk_datasets, "resolve_dataset_reference", resolve) + sources = Mock(return_value=[("enhanced_frs_2024_25", {"uri": "hf://x/y/z"})]) + monkeypatch.setattr(uk_datasets, "year_file_sources", sources) assert uk_datasets.load_datasets(years=[]) == {} - resolve.assert_called_once_with("uk", "enhanced_frs_2024_25") + sources.assert_called_once_with("uk", ["enhanced_frs_2024_25"], True) def test_uk_ensure_datasets_defaults_to_certified_bundle_dataset(monkeypatch): diff --git a/tests/test_uk_year_file_stamp.py b/tests/test_uk_year_file_stamp.py new file mode 100644 index 00000000..2df7f03a --- /dev/null +++ b/tests/test_uk_year_file_stamp.py @@ -0,0 +1,1709 @@ +"""Unit and property tests for the stamp on UK year files. + +``create_datasets`` stamps each UK year file, after its tables and data year, +with the identity of the source dataset it was cut from, the +policyengine-core and policyengine-uk versions that projected it, the year it +was cut for and a year-file format revision +(``uk.datasets.YEAR_FILE_STAMP_H5_DATASET``). +``ensure_datasets`` cuts a year file again, and ``load_datasets`` refuses it, +unless it records its data year and its stamp matches the one a file cut now +from the requested dataset would carry. These tests state the check's +invariants over generated stamps, data-year nodes, sources and package +versions, without a country engine: + +- **Exactness.** A year file is current exactly when it records its data + year and its stamp has this release's format revision, the requested year + and source, and the installed package versions. +- **Totality.** Whatever the stamp and the data-year node 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 differences.** The reason names both years and both sources when + they differ, and every package whose version differs, and nothing that + agrees. +- **Country isolation.** A UK year file's stamp is stored under a name of its + own, so a US year-file stamp never passes as a UK one, nor the reverse. +- **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: it never opens a table, never reads, + downloads or hashes the source, and never imports policyengine-uk. +- **Differential.** The UK and US checks name a source or package + difference with the same clauses, from the shared + ``year_file_stamp.source_and_package_differences``. + +``test_uk_year_file_stamp_integration.py`` runs ``create_datasets`` and +``ensure_datasets`` over a real policyengine-uk engine, and +``test_dataset_source_identity.py`` tests the source identities themselves. +""" + +from __future__ import annotations + +import builtins +import functools +import io +import json +import os +import tempfile +from importlib import metadata as importlib_metadata +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 + +import policyengine.provenance.dataset_materialization as dataset_materialization +from policyengine.provenance.dataset_materialization import ( + DatasetMaterializationError, + describe_dataset_source_identity, + requested_dataset_source_identity, +) +from policyengine.provenance.year_file_stamp import NOT_HDF5 +from policyengine.tax_benefit_models.uk import datasets as uk_datasets +from policyengine.tax_benefit_models.uk.datasets import ( + DATA_YEAR_KEY, + NO_DATA_YEAR, + STAMPED_WITHOUT_DATA_YEAR, + YEAR_FILE_PACKAGES, + YEAR_FILE_STAMP_H5_DATASET, + YEAR_FILE_STAMP_REVISION, + PolicyEngineUKDataset, + UKYearData, + _records_data_year, + _year_file_records_data_year, + ensure_datasets, + load_datasets, + read_year_file_stamp, + write_year_file_stamp, + year_file_stamp, + year_file_stamp_problem, +) + +REVISION = YEAR_FILE_STAMP_REVISION +DATA_YEAR = 2024 +#: ``hf://`` references are used as they are, and name year files after the +#: file's stem, so no manifest or download is involved. Its identity is the +#: URI as given. +SOURCE = "hf://fixture/repo/tiny.h5" +IDENTITY = {"uri": SOURCE} +STEM = "tiny" +#: The installed versions every test in this module sees (see +#: ``_installed_packages``), so no expectation depends on the environment. +PACKAGES = {"policyengine-core": "3.33.0", "policyengine-uk": "2.126.12"} +#: The real reader, kept before ``_installed_packages`` replaces it. +REAL_INSTALLED_PACKAGE_VERSIONS = uk_datasets.installed_package_versions +MALFORMED = "its stamp of what it was cut from and with is malformed" +NO_STAMP = ( + "it has no stamp of the source dataset and package versions it was cut " + "from and with, so it was cut by an older policyengine.py or not written " + "completely" +) + +#: Small alphabets of source identities in each of their three shapes (see +#: ``dataset_source_identity``), so generated sources often agree. +URIS = st.sampled_from([SOURCE, f"{SOURCE}@r1", f"{SOURCE}@r2"]) +SOURCES = st.one_of( + st.builds(lambda uri: {"uri": uri}, URIS), + st.builds( + lambda uri, sha256: {"uri": uri, "sha256": sha256}, + URIS, + st.sampled_from(["a" * 64, "b" * 64]), + ), + st.builds( + lambda path, size, mtime_ns: { + "path": path, + "size": size, + "mtime_ns": mtime_ns, + }, + st.sampled_from(["/data/tiny.h5", "/elsewhere/tiny.h5"]), + st.sampled_from([10, 11]), + # Before the epoch, at it, and past the year 9999. + st.sampled_from([-1, 0, 1_760_000_000_123_456_789, 10**20]), + ), +) +PACKAGE_NAMES = st.sampled_from([*YEAR_FILE_PACKAGES, "policyengine-us", "extra"]) +VERSIONS = st.one_of(st.none(), st.sampled_from(["1.0", "1.1"])) +PACKAGE_MAPS = st.dictionaries(PACKAGE_NAMES, VERSIONS, 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, +) + + +@pytest.fixture(autouse=True, scope="module") +def _installed_packages(): + """Pin the installed package versions a stamp records to ``PACKAGES``. + + Module-scoped, so Hypothesis tests can run under it. + """ + with pytest.MonkeyPatch.context() as patcher: + patcher.setattr( + uk_datasets, "installed_package_versions", lambda: dict(PACKAGES) + ) + yield + + +def _frames() -> dict[str, pd.DataFrame]: + return { + "person": pd.DataFrame( + { + "person_id": [1, 2], + "person_benunit_id": [1, 1], + "person_household_id": [1, 1], + "person_weight": [1.0, 1.0], + "age": [70, 68], + } + ), + "benunit": pd.DataFrame({"benunit_id": [1], "benunit_weight": [1.0]}), + "household": pd.DataFrame( + {"household_id": [1], "household_weight": [1.0], "region": ["LONDON"]} + ), + } + + +def _year_dataset(path, data_year=DATA_YEAR, year=DATA_YEAR) -> PolicyEngineUKDataset: + """A one-household year file for ``year``, projected from ``data_year`` + (or observed in ``year`` when ``data_year`` is ``None``).""" + tables = UKYearData( + **{ + entity: MicroDataFrame(frame, weights=f"{entity}_weight") + for entity, frame in _frames().items() + } + ) + projected = data_year is not None and year > data_year + return PolicyEngineUKDataset( + name="year-file", + description="Uncertified two-person UK year-file fixture", + filepath=str(path), + year=year, + data=tables, + data_year=data_year, + data_year_data=tables if projected else None, + ) + + +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") + ) + + +#: How a year file's data-year node can be stored: ``True`` as pandas writes +#: it (a group), ``False`` for none, or a node that is not a data year. +DATA_YEAR_NODES = ("dataset", "dangling link", "link cycle", "long link chain") + + +def _write_data_year_node(file, kind) -> None: + if kind == "dataset": + file.create_dataset(DATA_YEAR_KEY, data=DATA_YEAR) + elif kind == "dangling link": + file[DATA_YEAR_KEY] = h5py.SoftLink("/nowhere") + elif kind == "link cycle": + _link_cycle(file, DATA_YEAR_KEY) + elif kind == "long link chain": + # A chain too long to follow, to a group that would be a data year. + file.create_group(f"{DATA_YEAR_KEY}_end") + _link_chain(file, DATA_YEAR_KEY) + else: + raise AssertionError(kind) + + +def _write_year_file(path, *, stamp=_ABSENT, data_year=True, year=DATA_YEAR): + """Write a year file's tables and data year, then its stamp, as given. + + ``stamp`` is a JSON value, ``_ABSENT``, or a raw node (see + :func:`_write_node`). ``data_year`` is ``True`` for the data year as + ``save()`` writes it, ``False`` for none, or one of + :data:`DATA_YEAR_NODES`. A file for a ``year`` after the data year holds + the data year's tables too, as ``create_datasets`` writes it. + """ + Path(path).parent.mkdir(parents=True, exist_ok=True) + _year_dataset(path, DATA_YEAR if data_year is True else None, year).save() + with h5py.File(path, "a") as file: + if data_year not in (True, False): + _write_data_year_node(file, data_year) + _write_node(file, YEAR_FILE_STAMP_H5_DATASET, stamp) + return Path(path) + + +def _link_cycle(file, name) -> None: + """Point ``name`` into a cycle of soft links.""" + file[name] = h5py.SoftLink(f"/{name}_loop") + file[f"{name}_loop"] = h5py.SoftLink(f"/{name}") + + +def _link_chain(file, name) -> None: + """Point ``name`` down a chain of soft links too long to follow, to + ``{name}_end``.""" + file[name] = h5py.SoftLink(f"/{name}_0") + for index in range(100): + file[f"{name}_{index}"] = h5py.SoftLink(f"/{name}_{index + 1}") + file[f"{name}_100"] = h5py.SoftLink(f"/{name}_end") + + +#: What :func:`_stored` gives for a node no reader can parse as JSON. +_UNREADABLE = object() + + +def _stored(content): + """Decode what :func:`_write_node` stores, as the checker reads it back. + + An oracle independent of the checker's reader: raw text is parsed as + JSON; raw bytes are fixed-length strings, which read back without + trailing NULs; any other raw node, or text that is not JSON, is + unreadable. A JSON value reads back as its JSON round trip. + """ + if content is _ABSENT: + return _ABSENT + if isinstance(content, tuple): + kind, value = content + try: + if kind == "text": + return json.loads(value) + if kind == "bytes": + return json.loads(value.rstrip(b"\x00").decode("utf-8")) + except (RecursionError, ValueError): + pass + return _UNREADABLE + return json.loads(json.dumps(content)) + + +def _same_json(left, right) -> bool: + """Equal as JSON: ``True`` is not ``1``, and ``1.0`` is not ``1``.""" + return json.dumps(left, sort_keys=True) == json.dumps(right, sort_keys=True) + + +def _is_current( + stamp, data_year, source=IDENTITY, packages=PACKAGES, year=DATA_YEAR +) -> bool: + """The exactness invariant, stated directly on the stored values. + + Only a data year written as pandas writes it (a group) counts. + """ + return ( + data_year is True + and isinstance(stamp, dict) + and set(stamp) == {"packages", "revision", "source", "year"} + and type(stamp["revision"]) is int + and stamp["revision"] == REVISION + and type(stamp["year"]) is int + and stamp["year"] == year + and _same_json(stamp["source"], source) + and _same_json(stamp["packages"], packages) + ) + + +def _stamp(revision=REVISION, source=IDENTITY, packages=PACKAGES, year=DATA_YEAR): + return { + "packages": dict(packages), + "revision": revision, + "source": dict(source), + "year": year, + } + + +def _current(year=DATA_YEAR): + return year_file_stamp(IDENTITY, year) + + +# --- The stamp ---------------------------------------------------------- + + +def test_the_stamp_holds_the_source_packages_revision_and_year(): + assert _current() == { + "packages": PACKAGES, + "revision": 1, + "source": IDENTITY, + "year": DATA_YEAR, + } + assert year_file_stamp(IDENTITY, 2030)["year"] == 2030 + + +def test_the_stamp_reads_the_packages_when_called(monkeypatch): + """datasets.py must not freeze the versions at import.""" + monkeypatch.setattr( + uk_datasets, "installed_package_versions", lambda: {"z": None, "y": "1"} + ) + + assert year_file_stamp(IDENTITY, DATA_YEAR) == { + "packages": {"y": "1", "z": None}, + "revision": 1, + "source": IDENTITY, + "year": DATA_YEAR, + } + + +def test_explicit_packages_override_the_installed_ones(): + packages = {"policyengine-uk": "9.9.9"} + assert year_file_stamp(IDENTITY, DATA_YEAR, packages)["packages"] == packages + + +def test_the_stamp_is_stored_as_utf8_json(tmp_path): + path = _write_year_file(tmp_path / "year.h5") + write_year_file_stamp(path, _current()) + + with h5py.File(path, "r") as file: + node = file[YEAR_FILE_STAMP_H5_DATASET] + assert node.shape == () + assert json.loads(node.asstr()[()]) == _current() + + +def test_installed_package_versions_reads_metadata(monkeypatch): + """The real reader: distribution metadata, ``None`` when absent.""" + monkeypatch.setattr( + uk_datasets, + "YEAR_FILE_PACKAGES", + ("pytest", "hypothesis", "not-a-distribution-policyengine-test"), + ) + + assert REAL_INSTALLED_PACKAGE_VERSIONS() == { + "pytest": importlib_metadata.version("pytest"), + "hypothesis": importlib_metadata.version("hypothesis"), + "not-a-distribution-policyengine-test": None, + } + + +def test_the_stamp_records_policyengine_core_and_policyengine_uk(): + """policyengine-uk projects each year's tables, uprating by its parameters + through policyengine-core's parameter tree.""" + assert YEAR_FILE_PACKAGES == ("policyengine-core", "policyengine-uk") + assert set(REAL_INSTALLED_PACKAGE_VERSIONS()) == set(YEAR_FILE_PACKAGES) + + +STAMPS = st.fixed_dictionaries( + { + "packages": PACKAGE_MAPS, + "revision": st.integers(), + "source": SOURCES, + "year": st.integers(), + } +) + + +@settings(max_examples=50, deadline=None) +@given(first=STAMPS, second=STAMPS) +def test_a_year_file_stamp_reads_back_as_written(first, second): + with tempfile.TemporaryDirectory() as directory: + path = _write_year_file(Path(directory) / "year.h5") + write_year_file_stamp(path, first) + assert read_year_file_stamp(path) == first + # Writing again replaces it. + write_year_file_stamp(path, second) + assert read_year_file_stamp(path) == second + + +def test_a_missing_or_foreign_file_has_no_stamp(tmp_path): + assert read_year_file_stamp(tmp_path / "missing.h5") is None + (tmp_path / "text.h5").write_text("not an H5 file") + assert read_year_file_stamp(tmp_path / "text.h5") is None + assert read_year_file_stamp(_write_year_file(tmp_path / "year.h5")) is None + + +def test_the_stamp_does_not_disturb_pandas_or_loading(tmp_path): + path = _write_year_file(tmp_path / "year.h5", stamp=_current()) + + with pd.HDFStore(path, mode="r") as store: + assert sorted(store.keys()) == sorted( + [f"/{key}" for key in _frames()] + [f"/{DATA_YEAR_KEY}"] + ) + loaded = PolicyEngineUKDataset( + name="year-file", description="reload", filepath=str(path), year=DATA_YEAR + ) + assert loaded.data_year == DATA_YEAR + assert pd.DataFrame(loaded.data.person)["age"].tolist() == [70, 68] + + +# --- The data year ------------------------------------------------------ + + +@pytest.mark.parametrize("data_year", [True, False, "dataset"]) +def test_the_data_year_check_agrees_with_pandas(tmp_path, data_year): + """The check reads the data year's presence with h5py, as pandas lists + its keys, without opening the file with PyTables.""" + path = _write_year_file(tmp_path / "year.h5", data_year=data_year) + + with h5py.File(path, "r") as file: + assert _records_data_year(file) == _year_file_records_data_year(path) + assert _records_data_year(file) == (data_year is True) + + +def test_a_year_file_of_a_later_year_records_its_data_year(tmp_path): + """A projected file, with the data year's tables under their prefix.""" + tables = _year_dataset(tmp_path / "x.h5").data + path = tmp_path / "year.h5" + PolicyEngineUKDataset( + name="projected", + description="projected", + filepath=str(path), + year=DATA_YEAR + 2, + data=tables, + data_year=DATA_YEAR, + data_year_data=tables, + ).save() + write_year_file_stamp(path, _current()) + + assert year_file_stamp_problem(path, IDENTITY, DATA_YEAR) is None + + +def test_files_without_a_stamp_are_out_of_date_with_distinct_reasons(tmp_path): + """policyengine.py before #556 wrote neither; after it, the data year.""" + oldest = _write_year_file(tmp_path / "oldest.h5", data_year=False) + data_year_only = _write_year_file(tmp_path / "data_year_only.h5") + + assert year_file_stamp_problem(oldest, IDENTITY, DATA_YEAR) == NO_DATA_YEAR + assert year_file_stamp_problem(data_year_only, IDENTITY, DATA_YEAR) == NO_STAMP + assert NO_DATA_YEAR == ( + "it has no recorded data year, so it was written before policyengine.py " + "kept the observed year its tables were projected from, or not written " + "completely; policyengine-uk would take its projected tables as observed " + "data and uprate the State Pension by CPI rather than the triple lock" + ) + + +@pytest.mark.parametrize("data_year", [False, *DATA_YEAR_NODES]) +def test_a_stamped_file_without_a_data_year_is_out_of_date(tmp_path, data_year): + """A stamp alone is not enough: the run needs the data year. A data-year + node that is not a group, or a link that cannot be followed, is none.""" + path = _write_year_file(tmp_path / "year.h5", stamp=_current(), data_year=data_year) + + assert year_file_stamp_problem(path, IDENTITY, DATA_YEAR) == ( + STAMPED_WITHOUT_DATA_YEAR + ) + assert STAMPED_WITHOUT_DATA_YEAR == ( + "it is stamped but records no data year, which every year file " + "create_datasets writes records; policyengine-uk would take its projected " + "tables as observed data and uprate the State Pension by CPI rather than " + "the triple lock" + ) + + +# --- Exactness ---------------------------------------------------------- + + +@st.composite +def stamped_files(draw): + """A current source and packages, and a year file's stamp and data year + near them.""" + source = draw(SOURCES) + packages = draw(PACKAGE_MAPS) + stamped_source = draw(st.one_of(st.just(dict(source)), SOURCES)) + stamped_packages = draw( + st.one_of( + st.just(dict(packages)), + PACKAGE_MAPS, + # One version changed, dropped or added. + st.builds(lambda k, v: {**packages, k: v}, PACKAGE_NAMES, VERSIONS), + st.builds( + lambda name: {k: v for k, v in packages.items() if k != name}, + PACKAGE_NAMES, + ), + ) + ) + revision = draw(st.sampled_from([REVISION, REVISION, 0, 2, -1, 99])) + stamped_year = draw(st.sampled_from([DATA_YEAR, DATA_YEAR, DATA_YEAR + 1])) + stamp = draw( + st.one_of( + st.just(_stamp(revision, stamped_source, stamped_packages, stamped_year)), + st.just(_stamp(revision, stamped_source, stamped_packages, stamped_year)), + st.just(_ABSENT), + ) + ) + data_year = draw(st.sampled_from([True, True, True, False, *DATA_YEAR_NODES])) + return source, packages, stamp, data_year + + +@settings(max_examples=300, deadline=None) +@given(case=stamped_files()) +def test_a_year_file_is_current_only_under_the_exact_stamp_requested(case): + source, packages, stamp, data_year = case + with tempfile.TemporaryDirectory() as directory: + path = _write_year_file( + Path(directory) / "year.h5", stamp=stamp, data_year=data_year + ) + problem = year_file_stamp_problem(path, source, DATA_YEAR, packages) + + assert (problem is None) == _is_current(_stored(stamp), data_year, source, packages) + assert problem is None or (isinstance(problem, str) and problem) + + +#: 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( + { + "packages": st.one_of(PACKAGE_MAPS, JSON_VALUES), + "revision": st.one_of( + st.just(REVISION), st.booleans(), st.floats(), JSON_VALUES + ), + "source": st.one_of( + SOURCES, + JSON_VALUES, + # Well-formed keys holding values of the wrong type. + st.fixed_dictionaries( + {"path": JSON_VALUES, "size": JSON_VALUES, "mtime_ns": JSON_VALUES} + ), + st.fixed_dictionaries({"uri": JSON_VALUES, "sha256": JSON_VALUES}), + ), + "year": st.one_of( + st.sampled_from([DATA_YEAR, DATA_YEAR + 1]), + 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=400, deadline=None) +@given( + stamp=NODES, + data_year=st.sampled_from([True, False, *DATA_YEAR_NODES]), + source=SOURCES, + packages=PACKAGE_MAPS, +) +def test_the_year_file_problem_is_total_and_exact_for_any_content( + stamp, data_year, source, packages +): + """Totality: any stamp and data-year node give a reason or ``None``, + never an error, and ``None`` exactly for the stamp requested, however it + is stored (``_stored`` decodes raw nodes independently of the checker).""" + with tempfile.TemporaryDirectory() as directory: + path = _write_year_file( + Path(directory) / "year.h5", stamp=stamp, data_year=data_year + ) + problem = year_file_stamp_problem(path, source, DATA_YEAR, packages) + + assert problem is None or (isinstance(problem, str) and problem) + assert (problem is None) == _is_current(_stored(stamp), data_year, source, packages) + + +@pytest.mark.parametrize( + "stamp", + [ + {**_stamp(), "revision": True}, + {**_stamp(), "revision": 1.0}, + {**_stamp(), "revision": "1"}, + {"revision": 1}, + {k: v for k, v in _stamp().items() if k != "source"}, + {k: v for k, v in _stamp().items() if k != "packages"}, + {k: v for k, v in _stamp().items() if k != "year"}, + {**_stamp(), "year": True}, + {**_stamp(), "year": float(DATA_YEAR)}, + {**_stamp(), "year": str(DATA_YEAR)}, + {**_stamp(), "year": None}, + {**_stamp(), "extra": None}, + # A US stamp's shape, with its register. + {**_stamp(), "register": {}}, + {**_stamp(), "packages": {"policyengine-uk": 2.126}}, + {**_stamp(), "packages": [["policyengine-uk", "2.126.12"]]}, + {**_stamp(), "source": SOURCE}, + {**_stamp(), "source": {}}, + {**_stamp(), "source": {"uri": 5}}, + {**_stamp(), "source": {"uri": SOURCE, "sha256": None}}, + {**_stamp(), "source": {"uri": SOURCE, "revision": "r1"}}, + {**_stamp(), "source": {"path": "/a", "size": True, "mtime_ns": 1}}, + {**_stamp(), "source": {"path": "/a", "size": 1, "mtime_ns": 1.0}}, + {**_stamp(), "source": {"path": "/a", "size": 1}}, + [_stamp()], + ("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) + + assert year_file_stamp_problem(path, IDENTITY, DATA_YEAR) == MALFORMED + + +@pytest.mark.parametrize("revision", [0, 2, 3, -1]) +def test_a_year_file_from_another_format_revision_is_out_of_date(tmp_path, revision): + # Another revision may shape the rest of its stamp, or keep its data year, + # differently, and is still reported by revision rather than as malformed + # or as lacking a data year. + stamp = {"revision": revision, "sources": [SOURCE]} + path = _write_year_file(tmp_path / "year.h5", stamp=stamp, data_year=False) + + assert year_file_stamp_problem(path, IDENTITY, DATA_YEAR) == ( + f"it was cut with year-file format revision {revision}, and this " + "version of policyengine.py reads only revision 1" + ) + + +def test_a_stamp_stored_as_a_fixed_length_string_is_read(tmp_path): + stamp = json.dumps(_current()).encode() + path = _write_year_file(tmp_path / "year.h5", stamp=("bytes", stamp)) + + assert year_file_stamp_problem(path, IDENTITY, DATA_YEAR) is None + assert read_year_file_stamp(path) == _current() + + +def test_raising_the_format_revision_makes_older_files_out_of_date( + tmp_path, monkeypatch +): + older = _write_year_file(tmp_path / "older.h5", stamp=_current()) + monkeypatch.setattr(uk_datasets, "YEAR_FILE_STAMP_REVISION", 7) + newer = _write_year_file(tmp_path / "newer.h5", stamp=_current()) + + assert _current()["revision"] == 7 + assert year_file_stamp_problem(newer, IDENTITY, DATA_YEAR) is None + assert year_file_stamp_problem(older, IDENTITY, DATA_YEAR) == ( + "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 (see the same constant in +#: ``test_us_year_file_stamp.py``). Either way the stamp is not well formed. +DEEP_JSON = "[" * 100_000 + "0" + "]" * 100_000 + + +def test_json_nested_too_deeply_is_out_of_date_not_an_error(tmp_path): + path = _write_year_file(tmp_path / "year.h5", stamp=("text", DEEP_JSON)) + + assert year_file_stamp_problem(path, IDENTITY, DATA_YEAR) == MALFORMED + # Unreadable (None) or parsed into nested lists, but never a stamp. + assert not isinstance(read_year_file_stamp(path), dict) + + +def test_a_recursion_error_while_parsing_is_unreadable_content(tmp_path, monkeypatch): + """The RecursionError path, on every platform (see ``DEEP_JSON``).""" + path = _write_year_file(tmp_path / "year.h5", stamp=_current()) + loads = json.loads + + def deep(text, *args, **kwargs): + with h5py.File(path, "r") as file: + if text == file[YEAR_FILE_STAMP_H5_DATASET].asstr()[()]: + raise RecursionError("maximum recursion depth exceeded") + return loads(text, *args, **kwargs) + + monkeypatch.setattr(json, "loads", deep) + + assert year_file_stamp_problem(path, IDENTITY, DATA_YEAR) == MALFORMED + + +def test_a_dangling_link_in_place_of_the_stamp_reads_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.ExternalLink( + str(tmp_path / "missing.h5"), "/stamp" + ) + + assert year_file_stamp_problem(path, IDENTITY, DATA_YEAR) == NO_STAMP + assert read_year_file_stamp(path) is None + + +@pytest.mark.parametrize("links", ["cycle", "chain"]) +def test_stamp_links_that_cannot_be_followed_are_unreadable_content(tmp_path, links): + """Totality: h5py raises RuntimeError for these, which reads as + unreadable content rather than escaping.""" + path = _write_year_file(tmp_path / "year.h5") + with h5py.File(path, "a") as file: + if links == "cycle": + _link_cycle(file, YEAR_FILE_STAMP_H5_DATASET) + else: + _write_node(file, f"{YEAR_FILE_STAMP_H5_DATASET}_end", _current()) + _link_chain(file, YEAR_FILE_STAMP_H5_DATASET) + + assert year_file_stamp_problem(path, IDENTITY, DATA_YEAR) == MALFORMED + assert read_year_file_stamp(path) is None + + +def test_a_file_stamped_for_the_requested_source_is_current(tmp_path): + path = _write_year_file(tmp_path / "year.h5", stamp=_current()) + + assert year_file_stamp_problem(path, IDENTITY, DATA_YEAR) is None + + +# --- Named differences -------------------------------------------------- + + +@settings(max_examples=300, deadline=None) +@given( + source=SOURCES, + stamped_source=SOURCES, + packages=PACKAGE_MAPS, + stamped_packages=PACKAGE_MAPS, + stamped_year=st.sampled_from([DATA_YEAR, DATA_YEAR, DATA_YEAR + 3]), +) +def test_a_stale_year_file_names_each_year_source_and_version_that_differs( + source, stamped_source, packages, stamped_packages, stamped_year +): + """The reason has a year clause first exactly when the years differ, a + source clause exactly when the sources differ, naming both, and one clause + for each package whose version differs or that only one side records, + naming it; nothing that agrees is mentioned.""" + with tempfile.TemporaryDirectory() as directory: + path = _write_year_file( + Path(directory) / "year.h5", + stamp=_stamp(REVISION, stamped_source, stamped_packages, stamped_year), + ) + problem = year_file_stamp_problem(path, source, DATA_YEAR, packages) + + differing = { + name + for name in stamped_packages.keys() | packages.keys() + if name not in stamped_packages + or name not in packages + or stamped_packages[name] != packages[name] + } + if stamped_source == source and not differing and stamped_year == DATA_YEAR: + assert problem is None + return + clauses = problem.split("; ") + year_clauses = [c for c in clauses if c.startswith("it was cut for ")] + if stamped_year == DATA_YEAR: + assert year_clauses == [] + else: + assert clauses[0] == f"it was cut for {stamped_year}, not {DATA_YEAR}" + assert year_clauses == clauses[:1] + clauses = [c for c in clauses if c not in year_clauses] + source_clauses = [c for c in clauses if c.startswith("it was cut from ")] + if stamped_source == source: + assert source_clauses == [] + else: + assert source_clauses == [ + f"it was cut from {describe_dataset_source_identity(stamped_source)}, " + "and the requested source is " + f"{describe_dataset_source_identity(source)}" + ] + package_clauses = [c for c in clauses if c not in source_clauses] + assert len(package_clauses) == len(differing) + for name, clause in zip(sorted(differing), package_clauses): + assert name in clause + for other in (stamped_packages.keys() | packages.keys()) - {name}: + assert other not in clause + + +@pytest.mark.parametrize( + "stamped_source, source, described", + [ + ( + { + "uri": "hf://policyengine/uk-data/enhanced_frs_2024_25.h5@1.58.0", + "sha256": "a" * 64, + }, + { + "uri": "hf://policyengine/uk-data/enhanced_frs_2024_25.h5@1.59.0", + "sha256": "b" * 64, + }, + "hf://policyengine/uk-data/enhanced_frs_2024_25.h5@1.58.0 (sha256 " + f"{'a' * 64}), and the requested source is " + "hf://policyengine/uk-data/enhanced_frs_2024_25.h5@1.59.0 (sha256 " + f"{'b' * 64})", + ), + ( + # A re-upload under the same revision name. + {"uri": "hf://org/repo/populace_uk_2023.h5@r", "sha256": "a" * 64}, + {"uri": "hf://org/repo/populace_uk_2023.h5@r", "sha256": "b" * 64}, + f"hf://org/repo/populace_uk_2023.h5@r (sha256 {'a' * 64}), and the " + "requested source is hf://org/repo/populace_uk_2023.h5@r (sha256 " + f"{'b' * 64})", + ), + ( + {"path": "/data/frs.h5", "size": 10, "mtime_ns": 1_760_000_000_000_000_001}, + {"path": "/data/frs.h5", "size": 12, "mtime_ns": 1_760_000_060_000_000_000}, + "/data/frs.h5 (10 bytes, modified 2025-10-09T08:53:20.000000001Z), " + "and the requested source is /data/frs.h5 (12 bytes, modified " + "2025-10-09T08:54:20.000000000Z)", + ), + ], +) +def test_a_year_file_cut_from_another_source_names_both( + tmp_path, stamped_source, source, described +): + path = _write_year_file(tmp_path / "year.h5", stamp=_stamp(source=stamped_source)) + + assert ( + year_file_stamp_problem(path, source, DATA_YEAR) + == f"it was cut from {described}" + ) + + +@pytest.mark.parametrize( + "stamped_packages, packages, reason", + [ + ( + {**PACKAGES, "policyengine-uk": "2.126.11"}, + PACKAGES, + "it was cut with policyengine-uk 2.126.11, and policyengine-uk " + "2.126.12 is installed", + ), + ( + PACKAGES, + {**PACKAGES, "policyengine-uk": None}, + "it was cut with policyengine-uk 2.126.12, and policyengine-uk is " + "not installed", + ), + ( + {"policyengine-uk": "2.126.12"}, + PACKAGES, + "it was cut with no recorded policyengine-core version, and " + "policyengine-core 3.33.0 is installed", + ), + ( + {**PACKAGES, "policyengine-us": "2.2.1"}, + PACKAGES, + "it records the policyengine-us version it was cut with, which this " + "version of policyengine.py does not check", + ), + ( + {"policyengine-core": "3.32.0", "policyengine-uk": "2.126.11"}, + PACKAGES, + "it was cut with policyengine-core 3.32.0, and policyengine-core " + "3.33.0 is installed; it was cut with policyengine-uk 2.126.11, and " + "policyengine-uk 2.126.12 is installed", + ), + ], +) +def test_a_year_file_cut_with_other_versions_names_them( + tmp_path, stamped_packages, packages, reason +): + path = _write_year_file( + tmp_path / "year.h5", stamp=_stamp(packages=stamped_packages) + ) + + assert year_file_stamp_problem(path, IDENTITY, DATA_YEAR, packages) == reason + + +def test_every_difference_is_named_in_order(tmp_path): + path = _write_year_file( + tmp_path / "year.h5", + stamp=_stamp( + source={"uri": f"{SOURCE}@old"}, + packages={**PACKAGES, "policyengine-uk": "2.126.11"}, + ), + ) + + assert year_file_stamp_problem(path, IDENTITY, DATA_YEAR) == ( + f"it was cut from {SOURCE}@old, and the requested source is {SOURCE}; " + "it was cut with policyengine-uk 2.126.11, and policyengine-uk " + "2.126.12 is installed" + ) + + +def test_a_year_file_cut_for_another_year_names_both(tmp_path): + path = _write_year_file(tmp_path / "year.h5", stamp=_current(DATA_YEAR + 1)) + + assert year_file_stamp_problem(path, IDENTITY, DATA_YEAR) == ( + f"it was cut for {DATA_YEAR + 1}, not {DATA_YEAR}" + ) + assert year_file_stamp_problem(path, IDENTITY, DATA_YEAR + 1) is None + + +def test_a_year_file_copied_to_another_years_name_is_cut_again(tmp_path): + """Its tables are the other year's, so it is not taken for this one.""" + data_folder = tmp_path / "data" + original = data_folder / f"{STEM}_year_2025.h5" + _write_year_file(original, stamp=_current(2025), year=2025) + copy = data_folder / f"{STEM}_year_2026.h5" + copy.write_bytes(original.read_bytes()) + + with pytest.raises(ValueError, match="it was cut for 2025, not 2026"): + load_datasets(datasets=[SOURCE], years=[2026], data_folder=str(data_folder)) + create = MagicMock(return_value="created") + with patch.object(uk_datasets, "create_datasets", create): + assert ( + ensure_datasets( + datasets=[SOURCE], years=[2025, 2026], data_folder=str(data_folder) + ) + == "created" + ) + + +# --- Country isolation -------------------------------------------------- + + +def test_a_us_year_file_stamp_never_passes_as_a_uk_one(tmp_path): + """Each country stores its stamp under its own name: a UK year file that + carries only a US stamp is unstamped for the UK check, and a UK stamp is + no stamp to the US reader.""" + from policyengine.tax_benefit_models.us import legacy_inputs + + assert YEAR_FILE_STAMP_H5_DATASET == "policyengine_uk_year_file_stamp" + assert legacy_inputs.YEAR_FILE_STAMP_H5_DATASET != YEAR_FILE_STAMP_H5_DATASET + us_stamped = _write_year_file(tmp_path / "us_stamped.h5") + legacy_inputs.write_year_file_stamp( + us_stamped, legacy_inputs.year_file_stamp(IDENTITY, PACKAGES) + ) + uk_stamped = _write_year_file(tmp_path / "uk_stamped.h5", stamp=_current()) + + assert year_file_stamp_problem(us_stamped, IDENTITY, DATA_YEAR) == NO_STAMP + assert read_year_file_stamp(us_stamped) is None + assert legacy_inputs.read_year_file_stamp(uk_stamped) is None + assert legacy_inputs.year_file_stamp_problem(uk_stamped, IDENTITY) is not None + assert year_file_stamp_problem(uk_stamped, IDENTITY, DATA_YEAR) is None + + +# --- Differential: the US check names differences the same way ----------- + + +@settings(max_examples=150, deadline=None) +@given( + source=SOURCES, + stamped_source=SOURCES, + packages=PACKAGE_MAPS, + stamped_packages=PACKAGE_MAPS, +) +def test_the_uk_and_us_checks_name_source_and_package_differences_alike( + source, stamped_source, packages, stamped_packages +): + """Two implementations of one semantics: a UK and a US year file that + differ from the request only in source and packages get the same + reason.""" + from policyengine.tax_benefit_models.us import legacy_inputs + from policyengine.tax_benefit_models.us.datasets import ( + PolicyEngineUSDataset, + USYearData, + ) + + groups = ("household", "tax_unit", "spm_unit", "family", "marital_unit") + person = pd.DataFrame( + { + "person_id": [1], + **{f"person_{group}_id": [1] for group in groups}, + "person_weight": [1.0], + } + ) + frames = {"person": person} | { + group: pd.DataFrame({f"{group}_id": [1], f"{group}_weight": [1.0]}) + for group in groups + } + with tempfile.TemporaryDirectory() as directory: + uk_path = _write_year_file( + Path(directory) / "uk.h5", + stamp=_stamp(REVISION, stamped_source, stamped_packages), + ) + us_path = Path(directory) / "us.h5" + PolicyEngineUSDataset( + name="us", + description="us", + filepath=str(us_path), + year=DATA_YEAR, + data=USYearData( + **{ + entity: MicroDataFrame(frame, weights=f"{entity}_weight") + for entity, frame in frames.items() + } + ), + metadata={legacy_inputs.RENAMES_RECORD_KEY: {}}, + ).save() + us_stamp = legacy_inputs.year_file_stamp(stamped_source, stamped_packages) + legacy_inputs.write_year_file_stamp(us_path, us_stamp) + + uk = year_file_stamp_problem(uk_path, source, DATA_YEAR, packages) + us = legacy_inputs.year_file_stamp_problem(us_path, source, packages) + + assert uk == us + + +# --- Partial files ------------------------------------------------------ + + +@functools.cache +def _current_year_file_bytes() -> bytes: + with tempfile.TemporaryDirectory() as directory: + path = _write_year_file(Path(directory) / "year.h5", stamp=_current()) + assert year_file_stamp_problem(path, IDENTITY, DATA_YEAR) 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, IDENTITY, DATA_YEAR) 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, IDENTITY, DATA_YEAR) == NOT_HDF5 + assert NOT_HDF5 == ( + "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=_current()) + + dataset = PolicyEngineUKDataset( + name="year-file", description="reload", filepath=str(path), year=DATA_YEAR + ) + dataset.save() + + # ``save()`` rewrites the tables and the data year it loaded, but not the + # stamp, which only ``create_datasets`` writes. + assert dataset.data_year == DATA_YEAR + assert read_year_file_stamp(path) is None + assert year_file_stamp_problem(path, IDENTITY, DATA_YEAR) == 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=_current()) + put = pd.HDFStore.put + + def put_until_household(self, key, value, *args, **kwargs): + if key == "household": + raise RuntimeError("stopped part-way") + return put(self, key, value, *args, **kwargs) + + monkeypatch.setattr(pd.HDFStore, "put", put_until_household) + with pytest.raises(RuntimeError, match="stopped part-way"): + _year_dataset(path).save() + + assert read_year_file_stamp(path) is None + # The data year is written after the tables, so it is missing too. + assert year_file_stamp_problem(path, IDENTITY, DATA_YEAR) == NO_DATA_YEAR + + +# --- Errors unrelated to the file's contents ---------------------------- + + +def _unreadable_by_permission(path) -> Path: + path = _write_year_file(path, stamp=_current()) + 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, IDENTITY, DATA_YEAR) + with pytest.raises(error): + load_datasets(datasets=[SOURCE], years=[2024], data_folder=str(data_folder)) + with patch.object(uk_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=_current()) + + 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, IDENTITY, DATA_YEAR) + + +ERRORS = [ + PermissionError(13, "Permission denied"), + IsADirectoryError(21, "Is a directory"), + BlockingIOError(35, "Unable to lock file"), +] + + +@pytest.mark.parametrize("error", ERRORS) +def test_errors_unrelated_to_the_contents_are_raised_while_reading_the_stamp( + tmp_path, monkeypatch, error +): + path = _write_year_file(tmp_path / "year.h5", stamp=_current()) + + def fail(self, *args, **kwargs): + raise error + + monkeypatch.setattr(h5py.Dataset, "asstr", fail) + with pytest.raises(type(error)): + year_file_stamp_problem(path, IDENTITY, DATA_YEAR) + + +@pytest.mark.parametrize("error", ERRORS) +def test_errors_unrelated_to_the_contents_are_raised_while_finding_the_data_year( + tmp_path, monkeypatch, error +): + path = _write_year_file(tmp_path / "year.h5", stamp=_current()) + get = h5py.Group.get + + def fail_for_the_data_year(self, name, *args, **kwargs): + if name == DATA_YEAR_KEY: + raise error + return get(self, name, *args, **kwargs) + + monkeypatch.setattr(h5py.Group, "get", fail_for_the_data_year) + with pytest.raises(type(error)): + year_file_stamp_problem(path, IDENTITY, DATA_YEAR) + + +# --- ensure_datasets and load_datasets ---------------------------------- + +#: Year-file states: what is written, and whether the file is current +#: (``None``: the file is missing). +STATES = { + "current": (dict(stamp=_current()), True), + "before #556": (dict(data_year=False), False), + "data year only": (dict(), False), + "stamped, no data year": (dict(stamp=_current(), data_year=False), False), + "older revision": (dict(stamp=_stamp(0)), False), + "newer revision": (dict(stamp={"revision": 2}), False), + "other source revision": ( + dict(stamp=_stamp(source={"uri": f"{SOURCE}@old"})), + False, + ), + "verified source": ( + dict(stamp=_stamp(source={"uri": SOURCE, "sha256": "a" * 64})), + False, + ), + "other policyengine-uk": ( + dict(stamp=_stamp(packages={**PACKAGES, "policyengine-uk": "2.126.11"})), + False, + ), + "other policyengine-core": ( + dict(stamp=_stamp(packages={**PACKAGES, "policyengine-core": "3.34.0"})), + False, + ), + "malformed stamp": (dict(stamp=("text", "not json")), False), + "data year link cycle": (dict(stamp=_current(), data_year="link cycle"), False), + # ``_for_year`` stamps it for a year ten years on from the file's. + "cut for another year": (dict(stamp=_stamp(year="another")), False), + "not HDF5": (None, False), + "missing": (None, None), +} + + +def _for_year(written: dict, year: int) -> dict: + """A state's contents for the year file of ``year``: a stamp that records + a year records this one.""" + stamp = written.get("stamp", _ABSENT) + if isinstance(stamp, dict) and "year" in stamp: + stamped = year + 10 if stamp["year"] == "another" else year + written = {**written, "stamp": {**stamp, "year": stamped}} + return written + + +def _make_state(path: Path, state: str, year: int) -> 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, **_for_year(written, year), year=year) + + +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=_current(year), + year=year, + ) + return {"created": True} + + +@settings(max_examples=80, 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, year) + + 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) + ) + + # The stale files, from the states alone, and the reason load_datasets + # gives for each when it is asked for that year only. + expected_warnings = [] + for year, state in zip(years, states): + if STATES[state][1] is not False: + continue + path = data_folder / f"{STEM}_year_{year}.h5" + with pytest.raises(ValueError) as refused: + load_datasets( + datasets=[SOURCE], years=[year], data_folder=str(data_folder) + ) + prefix = f"UK year file {path} is out of date: " + suffix = ". Regenerate it with ensure_datasets() or create_datasets()." + message = str(refused.value) + assert message.startswith(prefix) and message.endswith(suffix) + expected_warnings.append((str(path), message[len(prefix) : -len(suffix)])) + create = MagicMock(side_effect=_create_current_files) + with ( + patch.object(uk_datasets, "create_datasets", create), + patch.object(uk_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 + ) == sorted(expected_warnings) + for call in warning.call_args_list: + assert call.args[0] == ( + "UK 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] + for year in years: + assert loaded[f"{STEM}_{year}"].data_year == DATA_YEAR + + +@pytest.mark.parametrize("local_source", [False, True]) +@pytest.mark.parametrize( + "states", + [ + ["current", "current"], + ["current", "data year only"], + ["current", "before #556"], + ["current", "other source revision"], + ["current", "other policyengine-uk"], + ], +) +def test_ensure_decides_from_the_stamps_alone(tmp_path, states, local_source): + """Deciding never opens a year file's tables (only its stamp and whether + its data year exists), and never reads, hashes or downloads the source + dataset or imports policyengine-uk. A local source is only + ``stat()``-ed.""" + data_folder = tmp_path / "data" + years = [2024, 2025] + dataset, identity = SOURCE, IDENTITY + if local_source: + source_file = tmp_path / "sources" / f"{STEM}.h5" + source_file.parent.mkdir() + source_file.write_bytes(b"the source is never read") + dataset, identity = ( + str(source_file), + requested_dataset_source_identity("uk", str(source_file)), + ) + for year, state in zip(years, states): + written = _for_year(STATES[state][0], year) + stamp = written.get("stamp", _ABSENT) + if isinstance(stamp, dict) and stamp.get("source") == IDENTITY: + stamp = {**stamp, "source": identity} + _write_year_file( + data_folder / f"{STEM}_year_{year}.h5", + **{**written, "stamp": stamp}, + year=year, + ) + load, create = MagicMock(return_value="loaded"), MagicMock(return_value="created") + source_path = Path(dataset).resolve() if local_source else None + + 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 StampAndDataYearOnly: + """An H5 file that lets the decision read only the stamp and find the + data year.""" + + def __init__(self, path, mode="r", **kwargs): + assert mode == "r" + assert source_path is None or Path(path).resolve() != source_path + 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, DATA_YEAR_KEY), name + node = self._file.get(name, default) + if name == DATA_YEAR_KEY: + # Its presence only: the group's contents are not handed over. + return h5py.Group.__new__(h5py.Group) if node is not None else None + return node + + real_open, real_import = builtins.open, builtins.__import__ + + def open_anything_but_the_source(file, *args, **kwargs): + if ( + source_path is not None + and isinstance(file, str | os.PathLike) + and Path(file).resolve() == source_path + ): + refuse() + return real_open(file, *args, **kwargs) + + def import_anything_but_policyengine_uk(name, *args, **kwargs): + if name == "policyengine_uk" or name.startswith("policyengine_uk."): + raise AssertionError("deciding imported policyengine-uk") + return real_import(name, *args, **kwargs) + + with ( + patch.object(h5py, "File", StampAndDataYearOnly), + 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(builtins, "open", open_anything_but_the_source), + patch.object(io, "open", open_anything_but_the_source), + patch.object(builtins, "__import__", import_anything_but_policyengine_uk), + patch.object(dataset_materialization, "sha256_file", refuse), + patch.object(dataset_materialization.requests, "get", refuse), + patch.object(uk_datasets, "materialize_dataset", refuse), + patch.object(uk_datasets, "load_datasets", load), + patch.object(uk_datasets, "create_datasets", create), + ): + result = ensure_datasets( + datasets=[dataset], 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_why(tmp_path): + data_folder = tmp_path / "data" + path = data_folder / f"{STEM}_year_2024.h5" + _write_year_file( + path, stamp=_stamp(packages={**PACKAGES, "policyengine-uk": "2.126.11"}) + ) + + with pytest.raises(ValueError) as refused: + load_datasets(datasets=[SOURCE], years=[2024], data_folder=str(data_folder)) + + assert str(refused.value) == ( + f"UK year file {path} is out of date: it was cut with policyengine-uk " + "2.126.11, and policyengine-uk 2.126.12 is installed. Regenerate it with " + "ensure_datasets() or create_datasets()." + ) + + +def test_the_refusal_names_the_source_it_was_cut_from(tmp_path): + data_folder = tmp_path / "data" + path = data_folder / f"{STEM}_year_2024.h5" + _write_year_file(path, stamp=_stamp(source={"uri": f"{SOURCE}@old"})) + + with pytest.raises(ValueError) as refused: + load_datasets(datasets=[SOURCE], years=[2024], data_folder=str(data_folder)) + + assert str(refused.value) == ( + f"UK year file {path} is out of date: it was cut from {SOURCE}@old, and " + f"the requested source is {SOURCE}. Regenerate it with ensure_datasets() " + "or create_datasets()." + ) + + +def test_a_year_file_without_a_data_year_is_refused_as_before(tmp_path): + data_folder = tmp_path / "data" + path = data_folder / f"{STEM}_year_2024.h5" + _write_year_file(path, data_year=False) + + with pytest.raises(ValueError) as refused: + load_datasets(datasets=[SOURCE], years=[2024], data_folder=str(data_folder)) + + assert str(refused.value) == ( + f"UK year file {path} is out of date: {NO_DATA_YEAR}. Regenerate it " + "with ensure_datasets() or create_datasets()." + ) + + +def test_load_and_ensure_compare_with_the_installed_versions(tmp_path, monkeypatch): + """Both read the versions when called, not at import.""" + data_folder = tmp_path / "data" + _write_year_file(data_folder / f"{STEM}_year_2024.h5", stamp=_current()) + load_datasets(datasets=[SOURCE], years=[2024], data_folder=str(data_folder)) + + monkeypatch.setattr( + uk_datasets, + "installed_package_versions", + lambda: {**PACKAGES, "policyengine-uk": "2.127.0"}, + ) + with pytest.raises(ValueError, match="policyengine-uk 2.127.0 is installed"): + load_datasets(datasets=[SOURCE], years=[2024], data_folder=str(data_folder)) + create = MagicMock(return_value="created") + with patch.object(uk_datasets, "create_datasets", create): + assert ( + ensure_datasets( + datasets=[SOURCE], years=[2024], data_folder=str(data_folder) + ) + == "created" + ) + + +# --- Two sources never share year files -------------------------------- + +#: Requests whose stems (``pop`` or ``other``) collide in different ways. +COLLIDING = [ + "hf://fixture/repo/pop.h5@r1", + "hf://fixture/repo/pop.h5@r2", + "hf://elsewhere/repo/pop.h5@r1", + "hf://fixture/repo/other.h5@r1", +] + + +def _stem_of(uri: str) -> str: + return uri.rsplit("/", 1)[-1].split(".")[0] + + +@settings(max_examples=60, deadline=None) +@given(requests=st.lists(st.sampled_from(COLLIDING), min_size=1, max_size=4)) +def test_two_sources_with_one_stem_are_refused_before_anything_is_cut(requests): + """Requests that would share year files are refused by all three + functions, before reading a year file or cutting one; requests of one + source repeated, or of sources with different stems, are not.""" + stems = {} + for request in requests: + stems.setdefault(_stem_of(request), set()).add(request) + collide = any(len(sources) > 1 for sources in stems.values()) + create = MagicMock(side_effect=lambda *a, **k: {}) + + def refuse(*args, **kwargs): + raise AssertionError("a year file was read or cut") + + with tempfile.TemporaryDirectory() as directory: + with ( + patch.object(uk_datasets, "materialize_dataset", refuse), + patch.object(uk_datasets, "year_file_stamp_problem", refuse), + patch.object(uk_datasets, "PolicyEngineUKDataset", refuse), + ): + for call in ( + lambda: load_datasets( + datasets=requests, years=[2024], data_folder=directory + ), + lambda: uk_datasets.create_datasets( + datasets=requests, + years=[2024], + data_folder=directory, + allow_unmanaged=True, + ), + ): + if collide: + with pytest.raises(ValueError, match="would share the UK year"): + call() + with patch.object(uk_datasets, "create_datasets", create): + if collide: + with pytest.raises(ValueError, match="would share the UK year"): + ensure_datasets( + datasets=requests, years=[2024], data_folder=directory + ) + create.assert_not_called() + else: + ensure_datasets(datasets=requests, years=[2024], data_folder=directory) + create.assert_called_once() + + +def test_the_collision_refusal_names_both_requests_and_sources(tmp_path): + first, second = "hf://fixture/repo/pop.h5@r1", "hf://fixture/repo/pop.h5@r2" + + with pytest.raises(ValueError) as refused: + ensure_datasets( + datasets=[first, second], years=[2024], data_folder=str(tmp_path) + ) + + assert str(refused.value) == ( + f"Datasets {first!r} and {second!r} would share the UK year files " + f"pop_year_{{year}}.h5, but are different sources: {first} and {second}. " + "Request them in separate calls with separate data folders." + ) + + +def test_two_local_files_with_one_name_are_refused(tmp_path): + """The usual smoke-test collision: ``pop.h5`` in two folders.""" + first, second = tmp_path / "a" / "pop.h5", tmp_path / "b" / "pop.h5" + for path in (first, second): + path.parent.mkdir() + path.write_bytes(b"source") + + with pytest.raises(ValueError, match="would share the UK year files pop_year_"): + ensure_datasets( + datasets=[str(first), str(second)], + years=[2024], + data_folder=str(tmp_path / "data"), + allow_unmanaged=True, + ) + + +@pytest.mark.parametrize( + "requests, error", + [ + (["not_a_bundle_dataset", SOURCE], "Unknown managed dataset"), + ([SOURCE, "hf://fixture/repo/other.h5"], "Explicit dataset URIs bypass"), + ], +) +def test_create_refuses_what_materialize_would_with_its_error( + tmp_path, requests, error +): + """With several datasets, ``create_datasets`` checks them all for shared + year files first. A request ``materialize_dataset`` would refuse is + refused there with the same error, and nothing is downloaded.""" + + def refuse(*args, **kwargs): + raise AssertionError("a dataset was materialized or downloaded") + + with ( + patch.object(uk_datasets, "materialize_dataset", refuse), + patch("requests.get", refuse), + ): + with pytest.raises(ValueError, match=error): + uk_datasets.create_datasets( + datasets=requests, years=[2024], data_folder=str(tmp_path) + ) + + +@pytest.mark.parametrize( + "request_, error", + [ + ("s3://bucket/tiny.h5", "Unsupported explicit dataset URI"), + ("hf://owner/repo", "Invalid Hugging Face dataset URI"), + ], +) +def test_create_refuses_a_uri_materialize_would_before_cutting_anything( + tmp_path, request_, error +): + """Even with ``allow_unmanaged``, with ``materialize_dataset``'s error.""" + + def refuse(*args, **kwargs): + raise AssertionError("a dataset was materialized or downloaded") + + with patch.object(uk_datasets, "materialize_dataset", refuse): + with pytest.raises(DatasetMaterializationError, match=error): + uk_datasets.create_datasets( + datasets=[SOURCE, request_], + years=[2024], + data_folder=str(tmp_path), + allow_unmanaged=True, + ) + + +def test_one_source_requested_by_two_names_is_not_a_collision(tmp_path): + """The default's logical name and its own URI are one source.""" + from policyengine.provenance.manifest import get_release_manifest + + manifest = get_release_manifest("uk") + requests = [manifest.default_dataset, manifest.default_dataset_uri] + create = MagicMock(return_value="created") + + with patch.object(uk_datasets, "create_datasets", create): + assert ( + ensure_datasets(datasets=requests, years=[2024], data_folder=str(tmp_path)) + == "created" + ) + + +def test_the_default_is_checked_against_its_bundle_pin(tmp_path): + """``datasets=None`` checks the certified default, by its manifest entry.""" + from policyengine.provenance.manifest import get_release_manifest + + manifest = get_release_manifest("uk") + default = requested_dataset_source_identity("uk", None) + stem = manifest.default_dataset + data_folder = tmp_path / "data" + _write_year_file( + data_folder / f"{stem}_year_2026.h5", + stamp=_stamp(source=default, year=2026), + year=2026, + ) + load = MagicMock(return_value="loaded") + + with patch.object(uk_datasets, "load_datasets", load): + assert ensure_datasets(years=[2026], data_folder=str(data_folder)) == "loaded" + assert default["sha256"] == manifest.datasets[stem].sha256 + assert ( + year_file_stamp_problem( + data_folder / f"{stem}_year_2026.h5", + requested_dataset_source_identity("uk", manifest.default_dataset_uri), + 2026, + ) + is None + ) diff --git a/tests/test_uk_year_file_stamp_integration.py b/tests/test_uk_year_file_stamp_integration.py new file mode 100644 index 00000000..264c7b7d --- /dev/null +++ b/tests/test_uk_year_file_stamp_integration.py @@ -0,0 +1,551 @@ +"""UK year files are never reused after a re-pin, an edited source or an upgrade. + +These tests run ``create_datasets``, ``ensure_datasets`` and ``load_datasets`` +over a real policyengine-uk engine, on the seven-person 2024 dataset of +``test_uk_year_file_data_year.py``. Each replays one way a UK year file could +be reused although it holds data cut from something else: + +- the bundle re-pins a logical dataset name to another revision and sha256 + (as releases of ``enhanced_frs_2024_25`` or ``populace_uk_2023`` do); +- a local source is rewritten in place; +- policyengine-uk or policyengine-core is upgraded; +- the source is replaced while the year files are cut; +- another process cuts the same year files at the same time; +- year files were cut by an earlier release, before the stamp existed. + +In each case ``load_datasets`` refuses the old year files, naming what +changed, and ``ensure_datasets`` cuts them again from the source it resolves +to now, and then reuses them. ``test_uk_year_file_stamp.py`` states the +check's invariants without an engine. +""" + +from __future__ import annotations + +import os +import re +import stat +from pathlib import Path +from unittest.mock import patch + +import h5py +import pandas as pd +import pytest + +pytest.importorskip("policyengine_uk") + +import policyengine.provenance.dataset_materialization as dataset_materialization # noqa: E402 +import policyengine.provenance.manifest as manifest_module # noqa: E402 +from policyengine.provenance.dataset_materialization import ( # noqa: E402 + DatasetMaterializationError, + describe_dataset_source_identity, + requested_dataset_source_identity, +) +from policyengine.provenance.manifest import ( # noqa: E402 + ArtifactPathReference, + get_release_manifest, +) +from policyengine.tax_benefit_models.uk import datasets as uk_datasets # noqa: E402 +from policyengine.tax_benefit_models.uk.datasets import ( # noqa: E402 + YEAR_FILE_STAMP_H5_DATASET, + PolicyEngineUKDataset, + create_datasets, + ensure_datasets, + installed_package_versions, + load_datasets, + read_year_file_stamp, + year_file_stamp, + year_file_stamp_problem, +) +from policyengine.utils.hashing import sha256_file # noqa: E402 +from tests.test_uk_year_file_data_year import ( # noqa: E402 + DATA_YEAR, + _tables, + _write_source, +) + +YEAR = 2026 +#: Reported State Pension of the replacement source: every amount differs +#: from ``_tables()``'s. +REPLACED = (11_000.0, 9_500.0, 15_000.0, 12_500.0, 0.0, 0.0, 4_000.0) + + +def _replacement_tables() -> dict[str, pd.DataFrame]: + return _tables(reported=REPLACED) + + +def _reported(dataset) -> list[float]: + """The data year's reported State Pension, as stored in a year file.""" + observed = dataset.data_year_data if dataset.year > DATA_YEAR else dataset.data + return pd.DataFrame(observed.person)["state_pension_reported"].tolist() + + +def _ensure(source: str, data_folder: Path, years=(YEAR,)): + return ensure_datasets( + datasets=[source], + years=list(years), + data_folder=str(data_folder), + allow_unmanaged=True, + ) + + +def _assert_reused(source: str, data_folder: Path, years=(YEAR,)): + """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(uk_datasets, "create_datasets", refuse): + return _ensure(source, data_folder, years) + + +def _assert_cut_again(source: str, data_folder: Path, years=(YEAR,)): + """An out-of-date year file is cut again, once.""" + with patch.object( + uk_datasets, "create_datasets", wraps=uk_datasets.create_datasets + ) as create: + ensured = _ensure(source, data_folder, years) + create.assert_called_once() + return ensured + + +def _new_source(tmp_path: Path, name="tiny_uk_2024.h5", tables=None) -> str: + """Write the tiny 2024 dataset (or ``tables``) as a local source.""" + folder = tmp_path / "sources" + folder.mkdir(exist_ok=True) + return _write_source(folder / name, tables) + + +def _partial_files(data_folder: Path) -> list[Path]: + """Temporary year files left in a data folder.""" + return sorted(data_folder.glob(".*.partial")) + + +def _stored(path: Path, year: int = YEAR) -> PolicyEngineUKDataset: + """A year file read directly, without the stamp check.""" + return PolicyEngineUKDataset( + name="stored", description="stored", filepath=str(path), year=year + ) + + +def test_create_stamps_each_year_file_with_its_source_and_packages(tmp_path): + source = _new_source(tmp_path) + data_folder = tmp_path / "data" + + created = create_datasets( + datasets=[source], + years=[DATA_YEAR, YEAR], + data_folder=str(data_folder), + allow_unmanaged=True, + ) + + identity = requested_dataset_source_identity("uk", source) + status = Path(source).resolve().stat() + assert identity == { + "path": str(Path(source).resolve()), + "size": status.st_size, + "mtime_ns": status.st_mtime_ns, + } + for year in (DATA_YEAR, YEAR): + path = created[f"tiny_uk_2024_{year}"].filepath + assert read_year_file_stamp(path) == { + "packages": installed_package_versions(), + "revision": 1, + "source": identity, + "year": year, + } + assert read_year_file_stamp(path) == year_file_stamp(identity, year) + assert year_file_stamp_problem(path, identity, year) is None + # Each year file was written beside its path and moved into place, with + # the mode a file pandas created at its path would have. + assert _partial_files(data_folder) == [] + if os.name != "nt": + umask = os.umask(0) + os.umask(umask) + for year in (DATA_YEAR, YEAR): + mode = (data_folder / f"tiny_uk_2024_year_{year}.h5").stat().st_mode + assert stat.S_IMODE(mode) == 0o666 & ~umask + # The real versions, not a fixture's. + assert all(installed_package_versions().values()) + + reused = _assert_reused(source, data_folder, (DATA_YEAR, YEAR)) + assert reused[f"tiny_uk_2024_{YEAR}"].data_year == DATA_YEAR + assert _reported(reused[f"tiny_uk_2024_{YEAR}"]) == list( + _tables()["person"]["state_pension_reported"] + ) + + +def test_an_edited_local_source_cuts_its_year_files_again(tmp_path): + """The smoke-test workflow: rewrite a local dataset in place.""" + source = _new_source(tmp_path) + data_folder = tmp_path / "data" + first = _ensure(source, data_folder) + old = requested_dataset_source_identity("uk", source) + assert _reported(first[f"tiny_uk_2024_{YEAR}"]) == list( + _tables()["person"]["state_pension_reported"] + ) + + _write_source(Path(source), _replacement_tables()) + # A modification time of its own, whatever the file system's resolution. + os.utime(source, ns=(old["mtime_ns"] + 10**9, old["mtime_ns"] + 10**9)) + new = requested_dataset_source_identity("uk", source) + assert new != old + + path = data_folder / f"tiny_uk_2024_year_{YEAR}.h5" + reason = ( + f"it was cut from {describe_dataset_source_identity(old)}, and the " + f"requested source is {describe_dataset_source_identity(new)}" + ) + assert year_file_stamp_problem(path, new, YEAR) == reason + with pytest.raises(ValueError, match=re.escape(reason)): + load_datasets(datasets=[source], years=[YEAR], data_folder=str(data_folder)) + + second = _assert_cut_again(source, data_folder) + assert _reported(second[f"tiny_uk_2024_{YEAR}"]) == list(REPLACED) + assert read_year_file_stamp(path)["source"] == new + reloaded = _assert_reused(source, data_folder) + assert _reported(reloaded[f"tiny_uk_2024_{YEAR}"]) == list(REPLACED) + + +@pytest.mark.parametrize("package", ["policyengine-uk", "policyengine-core"]) +def test_a_package_upgrade_cuts_the_year_files_again(tmp_path, monkeypatch, package): + source = _new_source(tmp_path) + data_folder = tmp_path / "data" + _ensure(source, data_folder) + cut_with = installed_package_versions() + + monkeypatch.setattr( + uk_datasets, + "installed_package_versions", + lambda: {**cut_with, package: "999.0.0"}, + ) + reason = ( + f"it was cut with {package} {cut_with[package]}, and {package} 999.0.0 " + "is installed" + ) + with pytest.raises(ValueError, match=re.escape(reason)): + load_datasets(datasets=[source], years=[YEAR], data_folder=str(data_folder)) + + _assert_cut_again(source, data_folder) + path = data_folder / f"tiny_uk_2024_year_{YEAR}.h5" + assert read_year_file_stamp(path)["packages"][package] == "999.0.0" + _assert_reused(source, data_folder) + + +BUNDLE_NAME = "tiny_uk_bundle" +BUNDLE_URI = f"hf://fixture/tiny-uk-bundle/{BUNDLE_NAME}.h5" + + +def _pinned(revision: str, sha256: str): + """Patch the bundled UK manifest to pin ``BUNDLE_NAME`` as given, and + refuse every download: a verified copy is already in the data folder.""" + manifest = get_release_manifest("uk").model_copy(deep=True) + manifest.datasets[BUNDLE_NAME] = ArtifactPathReference( + path=f"{BUNDLE_NAME}.h5", + repo_id="fixture/tiny-uk-bundle", + repo_type="dataset", + revision=revision, + sha256=sha256, + ) + + def refuse(*args, **kwargs): + raise AssertionError("a bundle dataset was downloaded") + + return ( + patch.object(manifest_module, "get_release_manifest", lambda _: manifest), + patch.object( + dataset_materialization, "get_release_manifest", lambda _: manifest + ), + patch.object(dataset_materialization.requests, "get", refuse), + ) + + +def test_a_re_pinned_bundle_dataset_cuts_its_year_files_again(tmp_path): + """A release of a UK dataset in miniature: the bundle points a logical + name at a new revision and sha256, and year files cut from the old + release are not reused, although their names are unchanged.""" + data_folder = tmp_path / "data" + data_folder.mkdir() + # ``materialize_dataset`` keeps a verified copy at the artifact's name in + # the data folder, and downloads only when that copy's sha256 differs. + bundled = data_folder / f"{BUNDLE_NAME}.h5" + _write_source(bundled) + old_sha256 = sha256_file(bundled) + old = {"uri": f"{BUNDLE_URI}@release-1", "sha256": old_sha256} + path = data_folder / f"{BUNDLE_NAME}_year_{YEAR}.h5" + + first, second, third = _pinned("release-1", old_sha256) + with first, second, third: + ensured = ensure_datasets( + datasets=[BUNDLE_NAME], years=[YEAR], data_folder=str(data_folder) + ) + assert ensured[f"{BUNDLE_NAME}_{YEAR}"].filepath == str(path) + assert read_year_file_stamp(path)["source"] == old + assert requested_dataset_source_identity("uk", BUNDLE_NAME) == old + + def refuse(*args, **kwargs): + raise AssertionError("a current year file was cut again") + + with patch.object(uk_datasets, "create_datasets", refuse): + ensure_datasets( + datasets=[BUNDLE_NAME], years=[YEAR], data_folder=str(data_folder) + ) + + # The new release reports other State Pension amounts; its verified copy + # replaces the old one, as a download would. + _write_source(bundled, _replacement_tables()) + new_sha256 = sha256_file(bundled) + assert new_sha256 != old_sha256 + + first, second, third = _pinned("release-2", new_sha256) + with first, second, third: + reason = ( + f"it was cut from {BUNDLE_URI}@release-1 (sha256 {old_sha256}), and " + f"the requested source is {BUNDLE_URI}@release-2 (sha256 {new_sha256})" + ) + assert ( + year_file_stamp_problem( + path, requested_dataset_source_identity("uk", BUNDLE_NAME), YEAR + ) + == reason + ) + with pytest.raises(ValueError, match=re.escape(reason)): + load_datasets( + datasets=[BUNDLE_NAME], years=[YEAR], data_folder=str(data_folder) + ) + with patch.object( + uk_datasets, "create_datasets", wraps=uk_datasets.create_datasets + ) as create: + recut = ensure_datasets( + datasets=[BUNDLE_NAME], years=[YEAR], data_folder=str(data_folder) + ) + create.assert_called_once() + assert _reported(recut[f"{BUNDLE_NAME}_{YEAR}"]) == list(REPLACED) + assert read_year_file_stamp(path)["source"] == { + "uri": f"{BUNDLE_URI}@release-2", + "sha256": new_sha256, + } + + +def _replacing_the_source_after_the_first_save(source: str): + """Patch ``save()`` so the source is replaced, by another process in + effect, once the first year file is written.""" + original = Path(source).stat().st_mtime_ns + save = uk_datasets.PolicyEngineUKDataset.save + replaced = [] + + def save_then_replace_the_source(self): + save(self) + if not replaced: + replaced.append(True) + _write_source(Path(source), _replacement_tables()) + os.utime(source, ns=(original + 10**9, original + 10**9)) + + return patch.object( + uk_datasets.PolicyEngineUKDataset, "save", save_then_replace_the_source + ) + + +def test_a_source_replaced_while_cut_writes_no_year_file(tmp_path): + """Another process replaces the source while the year files are cut, so + no year file may claim it was cut from what is there now.""" + source = _new_source(tmp_path) + data_folder = tmp_path / "data" + + with _replacing_the_source_after_the_first_save(source): + with pytest.raises(DatasetMaterializationError) as raised: + create_datasets( + datasets=[source], + years=[DATA_YEAR, YEAR], + data_folder=str(data_folder), + allow_unmanaged=True, + ) + + assert "changed while they were cut" in str(raised.value) + assert "None of them was written" in str(raised.value) + for year in (DATA_YEAR, YEAR): + assert not (data_folder / f"tiny_uk_2024_year_{year}.h5").exists() + assert _partial_files(data_folder) == [] + # Cut again, from the replacement, and stamped with its identity. + recut = _assert_cut_again(source, data_folder, (DATA_YEAR, YEAR)) + assert _reported(recut[f"tiny_uk_2024_{YEAR}"]) == list(REPLACED) + assert read_year_file_stamp(data_folder / f"tiny_uk_2024_year_{YEAR}.h5")[ + "source" + ] == requested_dataset_source_identity("uk", source) + + +def test_a_failed_cut_leaves_the_year_files_already_there_whole(tmp_path): + """Year files cut before keep their own tables and stamp, which no longer + match the replaced source, so they are cut again.""" + source = _new_source(tmp_path) + data_folder = tmp_path / "data" + _ensure(source, data_folder) + path = data_folder / f"tiny_uk_2024_year_{YEAR}.h5" + before = read_year_file_stamp(path) + # Touch the source so the next ensure cuts again, then replace it mid-cut. + old = requested_dataset_source_identity("uk", source) + os.utime(source, ns=(old["mtime_ns"] + 10**9, old["mtime_ns"] + 10**9)) + + with _replacing_the_source_after_the_first_save(source): + with pytest.raises(DatasetMaterializationError): + _ensure(source, data_folder) + + assert read_year_file_stamp(path) == before + assert _reported(_stored(path)) == list( + _tables()["person"]["state_pension_reported"] + ) + assert _partial_files(data_folder) == [] + recut = _assert_cut_again(source, data_folder) + assert _reported(recut[f"tiny_uk_2024_{YEAR}"]) == list(REPLACED) + + +def test_a_cut_that_fails_part_way_leaves_no_temporary_file(tmp_path): + """Whatever the error (here a failed write of the second year file), the + temporary files go and the year files already there are untouched.""" + source = _new_source(tmp_path) + data_folder = tmp_path / "data" + _ensure(source, data_folder, (DATA_YEAR, YEAR)) + paths = [data_folder / f"tiny_uk_2024_year_{year}.h5" for year in (DATA_YEAR, YEAR)] + before = [path.read_bytes() for path in paths] + save = uk_datasets.PolicyEngineUKDataset.save + saved = [] + + def fail_on_the_second_save(self): + if saved: + raise OSError(28, "No space left on device") + saved.append(self.filepath) + save(self) + + with patch.object( + uk_datasets.PolicyEngineUKDataset, "save", fail_on_the_second_save + ): + with pytest.raises(OSError, match="No space left"): + create_datasets( + datasets=[source], + years=[DATA_YEAR, YEAR], + data_folder=str(data_folder), + allow_unmanaged=True, + ) + + assert len(saved) == 1 + assert _partial_files(data_folder) == [] + assert [path.read_bytes() for path in paths] == before + + +def test_a_year_given_as_a_float_is_cut_and_returned(tmp_path): + """The result key and file name use the year as given, as before.""" + source = _new_source(tmp_path) + data_folder = tmp_path / "data" + + created = create_datasets( + datasets=[source], + years=[float(YEAR)], + data_folder=str(data_folder), + allow_unmanaged=True, + ) + + dataset = created[f"tiny_uk_2024_{float(YEAR)}"] + assert dataset.filepath == f"{data_folder}/tiny_uk_2024_year_{float(YEAR)}.h5" + assert read_year_file_stamp(dataset.filepath)["year"] == YEAR + assert _partial_files(data_folder) == [] + # ensure and load find the same file, and take it as current. + reused = _assert_reused(source, data_folder, (float(YEAR),)) + assert reused[f"tiny_uk_2024_{float(YEAR)}"].filepath == dataset.filepath + assert reused[f"tiny_uk_2024_{float(YEAR)}"].year == YEAR + + +def test_two_sources_cut_in_one_call_are_each_stamped_with_their_own(tmp_path): + first = _new_source(tmp_path) + second = _new_source(tmp_path, "other_uk_2024.h5", _replacement_tables()) + data_folder = tmp_path / "data" + + created = create_datasets( + datasets=[first, second], + years=[DATA_YEAR, YEAR], + data_folder=str(data_folder), + allow_unmanaged=True, + ) + + for source, stem, reported in ( + (first, "tiny_uk_2024", list(_tables()["person"]["state_pension_reported"])), + (second, "other_uk_2024", list(REPLACED)), + ): + identity = requested_dataset_source_identity("uk", source) + for year in (DATA_YEAR, YEAR): + path = data_folder / f"{stem}_year_{year}.h5" + assert created[f"{stem}_{year}"].filepath == str(data_folder / path.name) + assert read_year_file_stamp(path) == year_file_stamp(identity, year) + assert _reported(_stored(path, year)) == reported + with patch.object(uk_datasets, "create_datasets", side_effect=AssertionError): + ensure_datasets( + datasets=[first, second], + years=[DATA_YEAR, YEAR], + data_folder=str(data_folder), + allow_unmanaged=True, + ) + + +def test_concurrent_cuts_never_mix_one_cuts_tables_with_anothers_stamp(tmp_path): + """Two processes cut the same year files from different sources (two + ``pop.h5`` files, one data folder) at once: a second cut runs to the end + while the first is between writing a year file and stamping it. The file + then holds one cut's tables with that cut's stamp.""" + first_folder, second_folder = tmp_path / "a", tmp_path / "b" + first_folder.mkdir() + second_folder.mkdir() + first = _write_source(first_folder / "pop.h5") + second = _write_source(second_folder / "pop.h5", _replacement_tables()) + data_folder = tmp_path / "data" + save = uk_datasets.PolicyEngineUKDataset.save + interleaved = [] + + def save_then_let_the_other_cut_run(self): + save(self) + if not interleaved: + interleaved.append(True) + create_datasets( + datasets=[second], + years=[YEAR], + data_folder=str(data_folder), + allow_unmanaged=True, + ) + + with patch.object( + uk_datasets.PolicyEngineUKDataset, "save", save_then_let_the_other_cut_run + ): + create_datasets( + datasets=[first], + years=[YEAR], + data_folder=str(data_folder), + allow_unmanaged=True, + ) + + assert interleaved + path = data_folder / f"pop_year_{YEAR}.h5" + stamp = read_year_file_stamp(path) + tables = { + requested_dataset_source_identity("uk", first)["path"]: list( + _tables()["person"]["state_pension_reported"] + ), + requested_dataset_source_identity("uk", second)["path"]: list(REPLACED), + } + assert _reported(_stored(path)) == tables[stamp["source"]["path"]] + # The first cut finished last, so its year file is the one in place. + assert stamp["source"] == requested_dataset_source_identity("uk", first) + assert _partial_files(data_folder) == [] + + +def test_year_files_from_an_earlier_release_are_cut_again_once(tmp_path): + """Year files cut before the stamp record their data year but no stamp.""" + source = _new_source(tmp_path) + data_folder = tmp_path / "data" + _ensure(source, data_folder, (DATA_YEAR, YEAR)) + for year in (DATA_YEAR, YEAR): + with h5py.File(data_folder / f"tiny_uk_2024_year_{year}.h5", "a") as file: + del file[YEAR_FILE_STAMP_H5_DATASET] + + with pytest.raises(ValueError, match="it has no stamp of the source dataset"): + load_datasets( + datasets=[source], years=[DATA_YEAR, YEAR], data_folder=str(data_folder) + ) + _assert_cut_again(source, data_folder, (DATA_YEAR, YEAR)) + _assert_reused(source, data_folder, (DATA_YEAR, YEAR)) diff --git a/tests/test_us_year_file_stamp.py b/tests/test_us_year_file_stamp.py index bcd31cf8..0f162783 100644 --- a/tests/test_us_year_file_stamp.py +++ b/tests/test_us_year_file_stamp.py @@ -787,7 +787,7 @@ def deep(text, *args, **kwargs): raise RecursionError("maximum recursion depth exceeded") return loads(text, *args, **kwargs) - monkeypatch.setattr(legacy_inputs.json, "loads", deep) + monkeypatch.setattr(json, "loads", deep) assert year_file_stamp_problem(path, IDENTITY) == ( MALFORMED diff --git a/tests/test_year_file_stamp.py b/tests/test_year_file_stamp.py new file mode 100644 index 00000000..87f14514 --- /dev/null +++ b/tests/test_year_file_stamp.py @@ -0,0 +1,249 @@ +"""Property tests for the year-file stamp code the countries share. + +``policyengine.provenance.year_file_stamp`` holds what the US and UK +year-file checks have in common. ``test_us_year_file_stamp.py`` and +``test_uk_year_file_stamp.py`` test each country's check through it; these +tests state the shared pieces' own invariants: + +- **Reasons are exact.** ``source_and_package_differences`` gives no clause + exactly when the sources and package maps are equal, one clause when the + sources differ, and one per package that differs or only one side records. +- **Versions come from metadata.** ``installed_package_versions`` reads each + named distribution's version, ``None`` when it is not installed, and + imports nothing. +- **Stems never collide silently.** ``year_file_sources`` refuses two + requests that would share year files exactly when their stems agree and + their sources differ, naming the country. +- **Publishing is whole.** A year file written beside its path and published + replaces the file at the path in one step, with its stamp. +""" + +from __future__ import annotations + +import builtins +import os +from importlib import metadata as importlib_metadata +from pathlib import Path +from unittest.mock import patch + +import pytest +from hypothesis import given, settings +from hypothesis import strategies as st + +from policyengine.provenance.year_file_stamp import ( + installed_package_versions, + is_version_map, + publish_year_file, + read_stored_json, + source_and_package_differences, + temporary_year_file, + write_json_dataset, + year_file_sources, +) + +SOURCES = st.one_of( + st.builds( + lambda uri: {"uri": uri}, st.sampled_from(["hf://o/r/a.h5", "hf://o/r/b.h5"]) + ), + st.builds( + lambda sha256: {"uri": "hf://o/r/a.h5@v", "sha256": sha256}, + st.sampled_from(["a" * 64, "b" * 64]), + ), + st.builds( + lambda size: {"path": "/data/a.h5", "size": size, "mtime_ns": 0}, + st.sampled_from([1, 2]), + ), +) +PACKAGE_MAPS = st.dictionaries( + st.sampled_from(["policyengine-core", "policyengine-uk", "policyengine-us"]), + st.one_of(st.none(), st.sampled_from(["1.0", "1.1"])), + max_size=3, +) + + +@settings(max_examples=300, deadline=None) +@given( + source=SOURCES, + stamped_source=SOURCES, + packages=PACKAGE_MAPS, + stamped_packages=PACKAGE_MAPS, +) +def test_differences_are_named_exactly( + source, stamped_source, packages, stamped_packages +): + clauses = source_and_package_differences( + {"source": stamped_source, "packages": stamped_packages}, + {"source": source, "packages": packages}, + ) + + differing = { + name + for name in stamped_packages.keys() | packages.keys() + if stamped_packages.get(name, ...) != packages.get(name, ...) + } + assert (clauses == []) == (stamped_source == source and not differing) + assert len(clauses) == (stamped_source != source) + len(differing) + assert all(clause.startswith("it ") for clause in clauses) + package_clauses = clauses[1:] if stamped_source != source else clauses + for name, clause in zip(sorted(differing), package_clauses): + assert name in clause + + +def test_installed_package_versions_reads_metadata_without_importing(): + real_import = builtins.__import__ + + def refuse_policyengine(name, *args, **kwargs): + if name.startswith("policyengine_"): + raise AssertionError(f"imported {name}") + return real_import(name, *args, **kwargs) + + with patch.object(builtins, "__import__", refuse_policyengine): + versions = installed_package_versions( + ["pytest", "policyengine-core", "not-a-distribution-policyengine-test"] + ) + + assert versions == { + "pytest": importlib_metadata.version("pytest"), + "policyengine-core": importlib_metadata.version("policyengine-core"), + "not-a-distribution-policyengine-test": None, + } + assert is_version_map(versions) + + +@pytest.mark.parametrize( + "value, expected", + [ + ({}, True), + ({"a": "1", "b": None}, True), + ({"a": 1}, False), + ({1: "1"}, False), + ([["a", "1"]], False), + (None, False), + ], +) +def test_a_version_map_maps_names_to_versions_or_none(value, expected): + assert is_version_map(value) is expected + + +#: Requests by stem: ``pop`` from two revisions and two repositories, and one +#: ``other``. +REQUESTS = [ + "hf://fixture/repo/pop.h5@r1", + "hf://fixture/repo/pop.h5@r2", + "hf://elsewhere/repo/pop.h5@r1", + "hf://fixture/repo/other.h5@r1", +] + + +@settings(max_examples=100, deadline=None) +@given( + country=st.sampled_from(["us", "uk"]), + requests=st.lists(st.sampled_from(REQUESTS), min_size=1, max_size=4), +) +def test_year_file_sources_refuse_exactly_the_requests_that_collide(country, requests): + by_stem = {} + for request in requests: + by_stem.setdefault(request.rsplit("/", 1)[-1].split(".")[0], set()).add(request) + collide = any(len(sources) > 1 for sources in by_stem.values()) + + if collide: + with pytest.raises( + ValueError, match=f"would share the {country.upper()} year files" + ): + year_file_sources(country, requests) + return + sources = year_file_sources(country, requests) + assert sources == [ + (request.rsplit("/", 1)[-1].split(".")[0], {"uri": request}) + for request in requests + ] + + +def test_a_temporary_year_file_is_hidden_beside_its_path(tmp_path): + path = tmp_path / "data" / "pop_year_2026.h5" + + temporary = temporary_year_file(path) + + assert temporary.parent == path.parent + assert temporary.name.startswith(".pop_year_2026.") + assert temporary.name.endswith(".partial") + assert temporary.exists() and temporary.stat().st_size == 0 + assert temporary != temporary_year_file(path) + + +def _new_file_mode() -> int: + """The mode a newly created file gets: ``0o666`` less the umask.""" + import os + + umask = os.umask(0) + os.umask(umask) + return 0o666 & ~umask + + +@pytest.mark.skipif(os.name == "nt", reason="POSIX permission bits") +@pytest.mark.parametrize("umask", [0o022, 0o002, 0o077]) +def test_a_temporary_year_file_gets_a_new_files_mode(tmp_path, umask): + """Not ``mkstemp``'s owner-only 0600, so a year file moved into place is + as readable by others as one written at its path.""" + import os + import stat + + previous = os.umask(umask) + try: + temporary = temporary_year_file(tmp_path / "pop_year_2026.h5") + assert stat.S_IMODE(temporary.stat().st_mode) == _new_file_mode() + assert stat.S_IMODE(temporary.stat().st_mode) == 0o666 & ~umask + finally: + os.umask(previous) + + +def test_publishing_moves_a_stamped_file_onto_its_path_in_one_step(tmp_path): + """The file at the path is replaced whole: its old contents and stamp go + together, and the new contents arrive with their stamp.""" + import h5py + + path = tmp_path / "pop_year_2026.h5" + with h5py.File(path, "w") as file: + file["tables"] = [1] + write_json_dataset(path, "stamp", {"cut": "old"}) + temporary = temporary_year_file(path) + with h5py.File(temporary, "w") as file: + file["tables"] = [2] + + publish_year_file(temporary, path, "stamp", {"cut": "new"}) + + assert not temporary.exists() + assert read_stored_json(path, "stamp") == {"cut": "new"} + with h5py.File(path, "r") as file: + assert list(file["tables"][()]) == [2] + + +def test_stored_json_reads_none_for_absent_or_foreign_files(tmp_path): + assert read_stored_json(tmp_path / "missing.h5", "stamp") is None + (tmp_path / "text.h5").write_text("not an H5 file") + assert read_stored_json(tmp_path / "text.h5", "stamp") is None + write_json_dataset(tmp_path / "file.h5", "other", {"a": 1}) + assert read_stored_json(tmp_path / "file.h5", "stamp") is None + assert read_stored_json(tmp_path / "file.h5", "other") == {"a": 1} + + +def test_publishing_stamps_the_file_before_moving_it(tmp_path, monkeypatch): + """Stamped first, so a process that moves its own file onto the path in + between can never leave one cut's tables under another cut's stamp.""" + import os + + path = tmp_path / "pop_year_2026.h5" + temporary = temporary_year_file(path) + write_json_dataset(temporary, "tables", [1]) + replace = os.replace + moved = [] + + def replace_after_checking_the_stamp(source, destination): + assert read_stored_json(source, "stamp") == {"cut": "this"} + moved.append((Path(source), Path(destination))) + replace(source, destination) + + monkeypatch.setattr(os, "replace", replace_after_checking_the_stamp) + publish_year_file(temporary, path, "stamp", {"cut": "this"}) + + assert moved == [(temporary, path)]