diff --git a/.github/release-lock.md b/.github/release-lock.md new file mode 100644 index 00000000000..7797d104b58 --- /dev/null +++ b/.github/release-lock.md @@ -0,0 +1,55 @@ +# Release lock guard + +The automatic versioning job checks the committed `pyproject.toml` and `uv.lock` +before changing either file. It runs the existing `bump_version.py` exactly once, +builds the changelog with isolated PyPI tooling, and calls +`python .github/release_lock.py --refresh` before the `Update PolicyEngine US` +commit. That commit stages the version, lock, changelog, and consumed fragments. + +The guard validates the root package and every dependency source. Registry +packages must use `https://pypi.org/simple`; every listed sdist and wheel must +have a SHA-256 hash and an HTTPS artifact URL on `files.pythonhosted.org`. +A PyPI source label does not make a file, local path, or custom artifact URL +acceptable. Project source overrides, workspaces, and direct requirements are +also rejected, including ancestor workspaces that could redirect uv to another +lock. uv runs without inherited `UV_*`/`PIP_*` overrides or discovered +configuration and uses the standard PyPI registry. + +`--committed` requires both project and lock files to match their tracked HEAD +contents, validates their semantics, and runs an actual `uv lock --check`. +The default check supports inspecting an uncommitted candidate, but release CI +uses `--committed`. Neither check repairs a stale lock. + +`--refresh` is only for the automatic root version bump. It runs uv to refresh +and check the lock, preserving the complete reviewed dependency graph and all +lock metadata except the root package version. Any dependency, artifact, +constraint, marker, or other semantic change fails the step and restores the +exact previous lock bytes. A uv failure also restores those bytes. Solver +upgrades that change the graph fail closed and require a separate lock review. +The failed workflow cannot reach the automatic commit. + +Pull request model jobs depend on `ReleaseLock`. Sentinel push model jobs have +the same prerequisite, and `Publish` checks its own committed checkout again +before installation or build. Model environments use locked synchronization; +subsequent commands use that environment without synchronizing again. + +`ReleaseLockTests` has no dependency on a model job or the country lock check. +It uses only the Python standard library and uv, so it can test the guard while +the country dependency registry is incomplete: + +```sh +python -m unittest discover -s .github/tests -p test_release_lock.py -v +RELEASE_LOCK_REAL_UV=1 python -m unittest discover -s .github/tests -p test_release_lock.py -v +``` + +The opt-in probe creates an isolated, minimal project with a standard PyPI +dependency and exercises actual uv locking. CI enables it. These commands do not +install or import the country model. + +For the SPM integration, the committed country lock is intentionally still +blocked until `spm-calculator==1.0.0` is available on PyPI and a production +registry lock is regenerated and reviewed separately. A stale root version or +older calculator resolution must fail the full PR and publication gates. +Passing guard tests does not clear that release prerequisite. Do not use local +wheel links or edit lock fields to manufacture a passing lock; the automatic +root-version refresh is not a general dependency update command. diff --git a/.github/release_lock.py b/.github/release_lock.py new file mode 100644 index 00000000000..ad698f31fdc --- /dev/null +++ b/.github/release_lock.py @@ -0,0 +1,315 @@ +"""Validate a release lock, or refresh only its automatic version bump. + +Dependency changes belong in a reviewed registry lock before versioning runs. +This helper uses only the standard library and never imports the country model. +""" + +from __future__ import annotations + +import argparse +import copy +import os +from pathlib import Path +import re +import subprocess +import sys +import tomllib +from urllib.parse import unquote, urlsplit + + +REPO_ROOT = Path(__file__).resolve().parents[1] +PYPI = "https://pypi.org/simple" + + +def load_toml(path: Path) -> dict: + """Read project or lock metadata without importing its package.""" + return tomllib.loads(path.read_text(encoding="utf-8")) + + +def normalized_name(name: str) -> str: + return re.sub(r"[-_.]+", "-", name).lower() + + +def validate_registry_project(project: dict) -> None: + """Reject local, direct-URL, workspace and alternate-index requirements.""" + uv = project.get("tool", {}).get("uv", {}) + forbidden = { + "sources", + "index", + "workspace", + "find-links", + "index-url", + "extra-index-url", + "default-index", + "no-index", + } + if forbidden.intersection(uv): + raise ValueError( + "Release project must use PyPI without source/index/workspace overrides" + ) + + def validate_requirements(value): + if isinstance(value, str): + # Registry requirements may have version constraints, extras and + # environment markers, but never a URL or filesystem path. + requirement = value.split(";", 1)[0].strip() + if not re.match(r"^[A-Za-z0-9][A-Za-z0-9._-]*", requirement) or any( + char in requirement for char in "@/:\\" + ): + raise ValueError( + "Release requirements must use the registry, not URLs or paths" + ) + elif isinstance(value, list): + for item in value: + validate_requirements(item) + elif isinstance(value, dict): + for item in value.values(): + validate_requirements(item) + else: + raise ValueError("Invalid release requirement metadata") + + metadata = project["project"] + for requirements in ( + metadata.get("dependencies", []), + metadata.get("optional-dependencies", {}), + project.get("dependency-groups", {}), + project.get("build-system", {}).get("requires", []), + uv.get("constraint-dependencies", []), + uv.get("override-dependencies", []), + uv.get("build-constraint-dependencies", []), + ): + validate_requirements(requirements) + + +def validate_artifact(artifact: dict, package_name: str) -> None: + """Require a hashed PyPI-hosted artifact, even under a PyPI source label.""" + if not isinstance(artifact, dict) or set(artifact) - { + "url", + "hash", + "size", + "upload-time", + }: + raise ValueError(f"Invalid registry artifact fields for {package_name}") + url = artifact.get("url") + digest = artifact.get("hash") + if not isinstance(url, str) or any( + char.isspace() or ord(char) < 32 for char in url + ): + raise ValueError(f"Missing or invalid registry artifact URL for {package_name}") + parsed = urlsplit(url) + path = unquote(parsed.path) + if ( + parsed.scheme != "https" + or parsed.netloc != "files.pythonhosted.org" + or parsed.query + or parsed.fragment + or "?" in url + or "#" in url + or not path.startswith("/packages/") + or "\\" in path + or "%" in path + or any(part in {"", ".", ".."} for part in path.split("/")[1:]) + or path != parsed.path + ): + raise ValueError( + f"Artifact for {package_name} must have a canonical PyPI files URL" + ) + if not isinstance(digest, str) or not re.fullmatch( + r"sha256:[0-9a-fA-F]{64}", digest + ): + raise ValueError(f"Artifact for {package_name} must have a SHA-256 hash") + + +def validate_lock_requirements(value) -> None: + """Reject alternate transports in lock dependency edges and metadata too.""" + if isinstance(value, dict): + if {"url", "git", "path", "directory", "editable", "virtual"}.intersection( + value + ): + raise ValueError("Release lock requirement contains a non-registry source") + if "registry" in value and value["registry"] != PYPI: + raise ValueError("Release lock requirement must use the PyPI registry") + for item in value.values(): + validate_lock_requirements(item) + elif isinstance(value, list): + for item in value: + validate_lock_requirements(item) + + +def validate_registry_lock( + project: dict, lock: dict, *, allow_previous_version: bool = False +) -> None: + """The editable root is the only package allowed outside ordinary PyPI.""" + metadata = project["project"] + root_name = normalized_name(metadata["name"]) + roots = [] + for package in lock.get("package", []): + name = package["name"] + if normalized_name(name) == root_name: + roots.append(package) + if package.get("source") != {"editable": "."}: + raise ValueError( + "Release lock root must be the current editable checkout" + ) + if "sdist" in package or "wheels" in package: + raise ValueError( + "Editable release root must not contain registry artifacts" + ) + else: + if package.get("source") != {"registry": PYPI}: + raise ValueError( + f"Release lock dependency {name} must use the PyPI registry" + ) + artifacts = [] + if "sdist" in package: + artifacts.append(package["sdist"]) + wheels = package.get("wheels", []) + if not isinstance(wheels, list): + raise ValueError(f"Invalid registry wheel list for {name}") + artifacts.extend(wheels) + if not artifacts: + raise ValueError( + f"Release lock dependency {name} has no registry artifacts" + ) + for artifact in artifacts: + validate_artifact(artifact, name) + for key in ( + "dependencies", + "optional-dependencies", + "dev-dependencies", + "metadata", + ): + validate_lock_requirements(package.get(key, {})) + if len(roots) != 1: + raise ValueError("Release lock must contain exactly one root package") + if not allow_previous_version and roots[0].get("version") != metadata["version"]: + raise ValueError("Release lock root version differs from pyproject.toml") + # uv normalizes spacing in specifier lists; --check validates their meaning. + if re.sub(r"\s", "", lock.get("requires-python", "")) != re.sub( + r"\s", "", metadata["requires-python"] + ): + raise ValueError("Release lock Python range differs from pyproject.toml") + + +def without_root_version(lock: dict, root_name: str) -> dict: + """Retain the entire reviewed lock graph except the editable root version.""" + result = copy.deepcopy(lock) + for package in result["package"]: + if normalized_name(package["name"]) == normalized_name(root_name): + package.pop("version", None) + return result + + +def resolver_environment() -> dict[str, str]: + """Discard inherited resolver overrides and active project environments.""" + env = { + key: value + for key, value in os.environ.items() + if not key.startswith(("UV_", "PIP_")) + and key not in {"VIRTUAL_ENV", "CONDA_PREFIX", "PYTHONPATH", "PYTHONHOME"} + } + env["UV_FROZEN"] = "0" + return env + + +def require_committed_files(root: Path) -> None: + """A locally repaired lock or project must not stand in for reviewed HEAD.""" + for name in ("pyproject.toml", "uv.lock"): + path = root / name + result = subprocess.run( + ["git", "show", f"HEAD:{name}"], cwd=root, check=True, capture_output=True + ) + if path.is_symlink() or path.read_bytes() != result.stdout: + raise ValueError( + f"Release check requires committed {name}; working bytes differ from HEAD" + ) + + +def reject_parent_workspaces(root: Path) -> None: + """Prevent uv from checking or rewriting an ancestor's workspace lock.""" + # --no-config and --no-sources do not disable workspace discovery. Refuse + # workspace ancestors even when they claim to exclude this checkout. + for parent in root.resolve().parents: + project_path = parent / "pyproject.toml" + if project_path.is_file() and "workspace" in load_toml(project_path).get( + "tool", {} + ).get("uv", {}): + raise ValueError("Release checkout must not be inside a uv workspace") + + +def check_release_lock( + root: Path = REPO_ROOT, *, refresh: bool = False, committed: bool = False +) -> None: + """Check with uv, or transactionally refresh only the bumped root version.""" + if refresh and committed: + raise ValueError("Use committed checks before the bump and refresh after it") + reject_parent_workspaces(root) + if committed: + require_committed_files(root) + project = load_toml(root / "pyproject.toml") + validate_registry_project(project) + lock_path = root / "uv.lock" + if lock_path.is_symlink(): + raise ValueError("Release lock must be a regular checkout file") + before_bytes = lock_path.read_bytes() + before = tomllib.loads(before_bytes.decode("utf-8")) + validate_registry_lock(project, before, allow_previous_version=refresh) + command = [ + "uv", + "lock", + "--no-config", + "--no-sources", + "--default-index", + PYPI, + "--python", + sys.executable, + "--no-python-downloads", + "--no-cache", + ] + kwargs = {"cwd": root, "env": resolver_environment(), "check": True} + successful = False + try: + if refresh: + subprocess.run(command, **kwargs) + after = load_toml(lock_path) + validate_registry_lock(project, after) + name = project["project"]["name"] + if without_root_version(before, name) != without_root_version(after, name): + raise ValueError( + "Versioning changed the reviewed dependency graph; prepare and review a registry lock first" + ) + checked_bytes = lock_path.read_bytes() + subprocess.run([*command, "--check"], **kwargs) + if lock_path.read_bytes() != checked_bytes: + raise ValueError("uv lock --check unexpectedly changed the release lock") + successful = True + finally: + # Includes resolver errors, invalid TOML, drift, failed final checks and + # interruptions. A partial or deleted resolver output is never retained. + if not successful: + lock_path.write_bytes(before_bytes) + + +def main() -> int: + parser = argparse.ArgumentParser(description=__doc__) + mode = parser.add_mutually_exclusive_group() + mode.add_argument( + "--committed", + action="store_true", + help="Require lock and project bytes from HEAD", + ) + mode.add_argument( + "--refresh", + action="store_true", + help="Refresh only the root version after the automatic bump", + ) + args = parser.parse_args() + try: + check_release_lock(refresh=args.refresh, committed=args.committed) + except (ValueError, OSError, KeyError, subprocess.CalledProcessError) as exc: + parser.error(str(exc)) + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/.github/tests/fixtures/registry/README.md b/.github/tests/fixtures/registry/README.md new file mode 100644 index 00000000000..9ce6148b2ed --- /dev/null +++ b/.github/tests/fixtures/registry/README.md @@ -0,0 +1,18 @@ +This minimal graph was generated by uv 0.11.7 from the adjacent `pyproject.toml`, +with the standard PyPI index and no source or configuration overrides. The +`idna==3.10` artifact URLs and hashes are real registry metadata; the fixture +project is an editable local root. This fixture is unrelated to the country +package's pending production lock. + +Generation command: + +```sh +uv lock --no-config --no-sources --default-index https://pypi.org/simple \ + --cache-dir /tmp/release-lock-fixture-cache --python 3.13 --offline +``` + +The offline generation used copied existing PyPI cache metadata for idna, after +live registry access was unavailable on the development lane. The complete graph +was written by uv (two packages resolved); its TOML was not manually assembled. +The opt-in `RELEASE_LOCK_REAL_UV=1` test runs the installed uv with normal registry +access and validates a root-only refresh without depending on country packages. diff --git a/.github/tests/fixtures/registry/pyproject.toml b/.github/tests/fixtures/registry/pyproject.toml new file mode 100644 index 00000000000..1a5515816c7 --- /dev/null +++ b/.github/tests/fixtures/registry/pyproject.toml @@ -0,0 +1,9 @@ +[project] +name = "release-lock-fixture" +version = "1.0.0" +requires-python = ">=3.11" +dependencies = ["idna==3.10"] + +[build-system] +requires = ["hatchling"] +build-backend = "hatchling.build" diff --git a/.github/tests/fixtures/registry/uv.lock b/.github/tests/fixtures/registry/uv.lock new file mode 100644 index 00000000000..5ebbf09e2c0 --- /dev/null +++ b/.github/tests/fixtures/registry/uv.lock @@ -0,0 +1,23 @@ +version = 1 +revision = 3 +requires-python = ">=3.11" + +[[package]] +name = "idna" +version = "3.10" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/f1/70/7703c29685631f5a7590aa73f1f1d3fa9a380e654b86af429e0934a32f7d/idna-3.10.tar.gz", hash = "sha256:12f65c9b470abda6dc35cf8e63cc574b1c52b11df2c86030af0ac09b01b13ea9", size = 190490, upload-time = "2024-09-15T18:07:39.745Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/76/c6/c88e154df9c4e1a2a66ccf0005a88dfb2650c1dffb6f5ce603dfbd452ce3/idna-3.10-py3-none-any.whl", hash = "sha256:946d195a0d259cbba61165e88e65941f16e9b36ea6ddb97f00452bae8b1287d3", size = 70442, upload-time = "2024-09-15T18:07:37.964Z" }, +] + +[[package]] +name = "release-lock-fixture" +version = "1.0.0" +source = { editable = "." } +dependencies = [ + { name = "idna" }, +] + +[package.metadata] +requires-dist = [{ name = "idna", specifier = "==3.10" }] diff --git a/.github/tests/test_release_lock.py b/.github/tests/test_release_lock.py new file mode 100644 index 00000000000..75c7e0c2474 --- /dev/null +++ b/.github/tests/test_release_lock.py @@ -0,0 +1,480 @@ +"""Release-lock regressions runnable without importing the country model. + +Run: python -m unittest discover -s .github/tests -p test_release_lock.py -v +Set RELEASE_LOCK_REAL_UV=1 to exercise the installed uv against ordinary PyPI. +""" + +import copy +import importlib.util +import os +from pathlib import Path +import shutil +import subprocess +import tempfile +import tomllib +import unittest +from unittest.mock import patch + + +HERE = Path(__file__).resolve().parent +FIXTURE = HERE / "fixtures" / "registry" +SPEC = importlib.util.spec_from_file_location( + "release_lock", HERE.parent / "release_lock.py" +) +release_lock = importlib.util.module_from_spec(SPEC) +SPEC.loader.exec_module(release_lock) + + +class RegistryFixture(unittest.TestCase): + def setUp(self): + self.temporary = tempfile.TemporaryDirectory( + prefix="country-release-lock-test-" + ) + self.addCleanup(self.temporary.cleanup) + self.root = Path(self.temporary.name) + for name in ("pyproject.toml", "uv.lock"): + shutil.copyfile(FIXTURE / name, self.root / name) + self.project = tomllib.loads((self.root / "pyproject.toml").read_text()) + self.lock = tomllib.loads((self.root / "uv.lock").read_text()) + self.original = (self.root / "uv.lock").read_bytes() + + def bump_project(self): + path = self.root / "pyproject.toml" + path.write_text( + path.read_text().replace('version = "1.0.0"', 'version = "1.0.1"') + ) + + def refreshed_bytes(self): + return self.original.replace(b'version = "1.0.0"', b'version = "1.0.1"') + + def assert_rejected(self, callback, *args, **kwargs): + with self.assertRaises( + (ValueError, RuntimeError, subprocess.CalledProcessError) + ): + callback(*args, **kwargs) + + +class RegistryValidationTests(RegistryFixture): + def test_genuine_uv_registry_graph_is_accepted(self): + release_lock.validate_registry_project(self.project) + release_lock.validate_registry_lock(self.project, self.lock) + + def test_stale_root_version_is_rejected_before_resolver_runs(self): + self.bump_project() + with patch.object(release_lock.subprocess, "run") as run: + self.assert_rejected(release_lock.check_release_lock, self.root) + run.assert_not_called() + + def test_root_source_must_be_only_current_editable_project(self): + for source in ( + {"editable": "../other"}, + {"virtual": "."}, + {"path": "."}, + {"editable": ".", "registry": "https://pypi.org/simple"}, + ): + with self.subTest(source=source): + lock = copy.deepcopy(self.lock) + lock["package"][1]["source"] = source + self.assert_rejected( + release_lock.validate_registry_lock, self.project, lock + ) + + def test_exactly_one_root_is_required_and_it_cannot_have_artifacts(self): + missing = copy.deepcopy(self.lock) + missing["package"].pop() + duplicate = copy.deepcopy(self.lock) + duplicate["package"].append(copy.deepcopy(duplicate["package"][1])) + artifact = copy.deepcopy(self.lock) + artifact["package"][1]["wheels"] = artifact["package"][0]["wheels"] + for lock in (missing, duplicate, artifact): + with self.subTest(lock=lock): + self.assert_rejected( + release_lock.validate_registry_lock, self.project, lock + ) + + def test_registry_label_cannot_hide_local_or_alternate_package_sources(self): + for source in ( + {"registry": "https://example.org/simple"}, + {"registry": "http://pypi.org/simple"}, + {"registry": "https://pypi.org/simple", "path": "/tmp/idna"}, + {"url": "https://files.pythonhosted.org/packages/idna.whl"}, + {"editable": "../idna"}, + {"path": "/tmp/idna"}, + ): + with self.subTest(source=source): + lock = copy.deepcopy(self.lock) + lock["package"][0]["source"] = source + self.assert_rejected( + release_lock.validate_registry_lock, self.project, lock + ) + + def test_every_sdist_and_wheel_requires_trusted_artifact_url(self): + urls = [ + "file:///tmp/idna.whl", + "/tmp/idna.whl", + "../idna.whl", + "https://example.org/packages/idna.whl", + "http://files.pythonhosted.org/packages/idna.whl", + "https://files.pythonhosted.org.evil.example/packages/idna.whl", + "https://user:pass@files.pythonhosted.org/packages/idna.whl", + "https://files.pythonhosted.org:8443/packages/idna.whl", + "https://files.pythonhosted.org:443/packages/idna.whl", + "https://files.pythonhosted.org/packages/idna.whl?download=1", + "https://files.pythonhosted.org/packages/idna.whl#fragment", + "https://files.pythonhosted.org/packages/", + "https://files.pythonhosted.org/not-packages/idna.whl", + "https://files.pythonhosted.org/packages/../idna.whl", + "https://files.pythonhosted.org/packages/%2e%2e/idna.whl", + "https://files.pythonhosted.org/packages//idna.whl", + "https://files.pythonhosted.org/packages/a\\idna.whl", + ] + for kind in ("sdist", "wheel"): + for url in urls: + with self.subTest(kind=kind, url=url): + lock = copy.deepcopy(self.lock) + package = lock["package"][0] + artifact = ( + package["sdist"] if kind == "sdist" else package["wheels"][0] + ) + artifact["url"] = url + self.assert_rejected( + release_lock.validate_registry_lock, self.project, lock + ) + + def test_every_artifact_requires_sha256_and_rejects_path_fields(self): + mutations = [ + {"hash": None}, + {"hash": "sha256:"}, + {"hash": "sha256:" + "0" * 63}, + {"hash": "sha256:" + "g" * 64}, + {"hash": "md5:" + "a" * 32}, + {"path": "/tmp/idna.whl"}, + {"url": None}, + ] + for kind in ("sdist", "wheel"): + for mutation in mutations: + with self.subTest(kind=kind, mutation=mutation): + lock = copy.deepcopy(self.lock) + package = lock["package"][0] + artifact = ( + package["sdist"] if kind == "sdist" else package["wheels"][0] + ) + for key, value in mutation.items(): + if value is None: + artifact.pop(key, None) + else: + artifact[key] = value + self.assert_rejected( + release_lock.validate_registry_lock, self.project, lock + ) + + def test_a_bad_later_wheel_is_also_rejected(self): + self.lock["package"][0]["wheels"].append( + {"url": "file:///tmp/second.whl", "hash": "sha256:" + "0" * 64} + ) + self.assert_rejected( + release_lock.validate_registry_lock, self.project, self.lock + ) + + def test_dependency_without_any_artifacts_is_rejected(self): + self.lock["package"][0].pop("sdist") + self.lock["package"][0].pop("wheels") + self.assert_rejected( + release_lock.validate_registry_lock, self.project, self.lock + ) + + def test_project_resolver_source_overrides_are_rejected(self): + entries = [ + ("sources", {"idna": {"path": "../idna"}}), + ("index", [{"url": "https://pypi.org/simple"}]), + ("workspace", {"members": ["packages/*"]}), + ("find-links", ["/tmp/wheels"]), + ("index-url", "https://pypi.org/simple"), + ("extra-index-url", ["https://example.org/simple"]), + ] + for key, value in entries: + with self.subTest(key=key): + project = copy.deepcopy(self.project) + project["tool"] = {"uv": {key: value}} + self.assert_rejected(release_lock.validate_registry_project, project) + + def test_direct_url_or_path_requirements_are_rejected_in_all_groups(self): + requirements = [ + "idna @ https://files.pythonhosted.org/packages/idna.whl", + "idna @ file:///tmp/idna", + "idna @ ../idna", + "../idna", + "/tmp/idna", + ] + for location in ("dependencies", "optional", "group", "build"): + for requirement in requirements: + with self.subTest(location=location, requirement=requirement): + project = copy.deepcopy(self.project) + if location == "dependencies": + project["project"]["dependencies"] = [requirement] + elif location == "optional": + project["project"]["optional-dependencies"] = { + "dev": [requirement] + } + elif location == "group": + project["dependency-groups"] = {"dev": [requirement]} + else: + project["build-system"]["requires"] = [requirement] + self.assert_rejected( + release_lock.validate_registry_project, project + ) + + def test_resolver_environment_discards_inherited_uv_and_pip_overrides(self): + inherited = { + "PATH": "/usr/bin", + "UV_INDEX": "https://example.org/simple", + "UV_FIND_LINKS": "/tmp/wheels", + "UV_OFFLINE": "1", + "UV_CONFIG_FILE": "/tmp/uv.toml", + "UV_WORKING_DIRECTORY": "/tmp", + "PIP_INDEX_URL": "https://example.org/simple", + "PIP_FIND_LINKS": "/tmp", + } + with patch.dict(os.environ, inherited, clear=True): + env = release_lock.resolver_environment() + self.assertEqual(env["PATH"], inherited["PATH"]) + for key in inherited: + if key.startswith(("UV_", "PIP_")): + self.assertNotIn(key, env) + + +class ReleaseTransactionTests(RegistryFixture): + def test_parent_workspace_cannot_redirect_the_lock_check(self): + child = self.root / "members" / "child" + child.mkdir(parents=True) + for name in ("pyproject.toml", "uv.lock"): + shutil.copyfile(FIXTURE / name, child / name) + (self.root / "pyproject.toml").write_text( + '[tool.uv.workspace]\nmembers = ["members/*"]\n' + ) + with patch.object(release_lock.subprocess, "run") as run: + self.assert_rejected(release_lock.check_release_lock, child) + run.assert_not_called() + self.assertEqual((child / "uv.lock").read_bytes(), self.original) + + def test_check_uses_uv_lock_check_without_writing_lock(self): + with patch.object(release_lock.subprocess, "run") as run: + release_lock.check_release_lock(self.root) + self.assertEqual((self.root / "uv.lock").read_bytes(), self.original) + command = run.call_args.args[0] + self.assertIn("lock", command) + self.assertIn("--check", command) + self.assertIn("--no-config", command) + self.assertIn("--no-sources", command) + self.assertIn("https://pypi.org/simple", command) + self.assertTrue(run.call_args.kwargs["check"]) + + def test_uv_check_failure_is_not_treated_as_success(self): + with patch.object( + release_lock.subprocess, + "run", + side_effect=subprocess.CalledProcessError(1, ["uv", "lock", "--check"]), + ): + self.assert_rejected(release_lock.check_release_lock, self.root) + self.assertEqual((self.root / "uv.lock").read_bytes(), self.original) + + def test_successful_check_that_changes_lock_is_rejected_and_restored(self): + def mutate(command, **kwargs): + (self.root / "uv.lock").write_bytes( + self.original + b"\n# unexpected write\n" + ) + return subprocess.CompletedProcess(command, 0) + + with patch.object(release_lock.subprocess, "run", side_effect=mutate): + self.assert_rejected(release_lock.check_release_lock, self.root) + self.assertEqual((self.root / "uv.lock").read_bytes(), self.original) + + def test_refresh_allows_only_root_version_change_then_checks(self): + self.bump_project() + commands = [] + + def resolve(command, **kwargs): + commands.append(command) + if "--check" not in command: + (self.root / "uv.lock").write_bytes(self.refreshed_bytes()) + return subprocess.CompletedProcess(command, 0) + + with patch.object(release_lock.subprocess, "run", side_effect=resolve): + release_lock.check_release_lock(self.root, refresh=True) + self.assertEqual(len(commands), 2) + self.assertNotIn("--check", commands[0]) + self.assertIn("--check", commands[1]) + for command in commands: + self.assertIn("--no-config", command) + self.assertIn("--no-sources", command) + self.assertIn("https://pypi.org/simple", command) + current = tomllib.loads((self.root / "uv.lock").read_text()) + self.assertEqual(current["package"][1]["version"], "1.0.1") + self.assertEqual( + release_lock.without_root_version(current, "release-lock-fixture"), + release_lock.without_root_version(self.lock, "release-lock-fixture"), + ) + + def test_any_reviewed_graph_drift_restores_previous_lock_bytes(self): + mutations = { + "dependency-version": lambda b: b.replace( + b'version = "3.10"', b'version = "3.11"' + ), + "artifact-hash": lambda b: b.replace(b"12f65c9b", b"02f65c9b"), + "root-requirement": lambda b: b.replace(b"==3.10", b">=3.10"), + "lock-python-range": lambda b: b.replace(b">=3.11", b">=3.12"), + "artifact-url": lambda b: b.replace( + b"files.pythonhosted.org", b"example.org" + ), + } + self.bump_project() + for name, mutate in mutations.items(): + with self.subTest(name=name): + + def resolve(command, **kwargs): + (self.root / "uv.lock").write_bytes(mutate(self.refreshed_bytes())) + return subprocess.CompletedProcess(command, 0) + + with patch.object(release_lock.subprocess, "run", side_effect=resolve): + self.assert_rejected( + release_lock.check_release_lock, self.root, refresh=True + ) + self.assertEqual((self.root / "uv.lock").read_bytes(), self.original) + + def test_network_failure_after_partial_write_restores_previous_lock_bytes(self): + self.bump_project() + + def fail(command, **kwargs): + (self.root / "uv.lock").write_bytes(b"partial resolver output\n") + raise subprocess.CalledProcessError( + 2, command, stderr="network unavailable" + ) + + with patch.object(release_lock.subprocess, "run", side_effect=fail): + self.assert_rejected( + release_lock.check_release_lock, self.root, refresh=True + ) + self.assertEqual((self.root / "uv.lock").read_bytes(), self.original) + + def test_interrupted_refresh_restores_a_deleted_previous_lock(self): + self.bump_project() + + def interrupt(command, **kwargs): + (self.root / "uv.lock").unlink() + raise KeyboardInterrupt() + + with patch.object(release_lock.subprocess, "run", side_effect=interrupt): + with self.assertRaises(KeyboardInterrupt): + release_lock.check_release_lock(self.root, refresh=True) + self.assertEqual((self.root / "uv.lock").read_bytes(), self.original) + + def test_post_refresh_check_failure_restores_previous_lock_bytes(self): + self.bump_project() + + def resolve(command, **kwargs): + if "--check" in command: + raise subprocess.CalledProcessError(1, command) + (self.root / "uv.lock").write_bytes(self.refreshed_bytes()) + return subprocess.CompletedProcess(command, 0) + + with patch.object(release_lock.subprocess, "run", side_effect=resolve): + self.assert_rejected( + release_lock.check_release_lock, self.root, refresh=True + ) + self.assertEqual((self.root / "uv.lock").read_bytes(), self.original) + + def test_invalid_previous_artifact_is_rejected_without_resolver(self): + invalid = self.original.replace( + b"https://files.pythonhosted.org", b"file:///tmp" + ) + (self.root / "uv.lock").write_bytes(invalid) + self.bump_project() + with patch.object(release_lock.subprocess, "run") as run: + self.assert_rejected( + release_lock.check_release_lock, self.root, refresh=True + ) + run.assert_not_called() + self.assertEqual((self.root / "uv.lock").read_bytes(), invalid) + + +class CommittedLockTests(RegistryFixture): + def setUp(self): + super().setUp() + self.git("init", "-q") + self.git("add", "pyproject.toml", "uv.lock") + self.git( + "-c", + "user.name=Release Lock Test", + "-c", + "user.email=release-lock@example.invalid", + "-c", + "commit.gpgsign=false", + "commit", + "-qm", + "Registry fixture", + ) + + def git(self, *args): + return subprocess.run( + ["git", *args], cwd=self.root, check=True, capture_output=True + ) + + def test_committed_check_accepts_exact_head_files(self): + real_run = subprocess.run + + def resolve(command, **kwargs): + if command[0] == "git": + return real_run(command, **kwargs) + return subprocess.CompletedProcess(command, 0) + + with patch.object(release_lock.subprocess, "run", side_effect=resolve): + release_lock.check_release_lock(self.root, committed=True) + + def test_committed_check_rejects_uncommitted_lock_and_project(self): + for filename in ("uv.lock", "pyproject.toml"): + with self.subTest(filename=filename): + path = self.root / filename + previous = path.read_bytes() + path.write_bytes(previous + b"\n# Uncommitted local regeneration\n") + try: + self.assert_rejected( + release_lock.check_release_lock, self.root, committed=True + ) + finally: + path.write_bytes(previous) + + +@unittest.skipUnless( + os.environ.get("RELEASE_LOCK_REAL_UV") == "1", + "set RELEASE_LOCK_REAL_UV=1 for the real registry resolver probe", +) +class RealUvTests(RegistryFixture): + def test_actual_registry_root_only_refresh(self): + """Resolve genuine PyPI idna, bump only root, and check the final graph.""" + env = release_lock.resolver_environment() + command = [ + "uv", + "lock", + "--no-config", + "--no-sources", + "--default-index", + "https://pypi.org/simple", + "--cache-dir", + str(self.root / "cache"), + ] + subprocess.run(command, cwd=self.root, env=env, check=True, timeout=60) + before = tomllib.loads((self.root / "uv.lock").read_text()) + release_lock.check_release_lock(self.root) + self.bump_project() + release_lock.check_release_lock(self.root, refresh=True) + after = tomllib.loads((self.root / "uv.lock").read_text()) + self.assertEqual(after["package"][1]["version"], "1.0.1") + self.assertEqual( + release_lock.without_root_version(before, "release-lock-fixture"), + release_lock.without_root_version(after, "release-lock-fixture"), + ) + release_lock.check_release_lock(self.root) + + +if __name__ == "__main__": + unittest.main() diff --git a/.github/workflows/pr.yaml b/.github/workflows/pr.yaml index 042dbecf8dd..6e149140457 100644 --- a/.github/workflows/pr.yaml +++ b/.github/workflows/pr.yaml @@ -7,16 +7,45 @@ concurrency: group: ${{ github.workflow }}-${{ github.ref }} cancel-in-progress: true jobs: + ReleaseLock: + name: Check committed registry lock + runs-on: ubuntu-latest + steps: + - uses: actions/checkout@v6 + - name: Set up Python + uses: actions/setup-python@v6 + with: + python-version: "3.14" + - name: Install uv + uses: astral-sh/setup-uv@v8.1.0 + - name: Check committed registry lock + run: python .github/release_lock.py --committed + ReleaseLockTests: + name: Release lock guard tests + runs-on: ubuntu-latest + steps: + - uses: actions/checkout@v6 + - name: Set up Python + uses: actions/setup-python@v6 + with: + python-version: "3.14" + - name: Install uv + uses: astral-sh/setup-uv@v8.1.0 + - name: Test guard without importing the country model + env: + RELEASE_LOCK_REAL_UV: "1" + run: python -m unittest discover -s .github/tests -p test_release_lock.py -v Lint: + needs: ReleaseLock runs-on: ubuntu-latest steps: - uses: actions/checkout@v6 - name: Install uv uses: astral-sh/setup-uv@v8.1.0 - name: Install dependencies - run: uv sync --extra dev + run: uv sync --locked --extra dev - name: Check formatting - run: uv run ruff format --check . + run: uv run --no-sync ruff format --check . check-changelog: name: Check changelog fragment runs-on: ubuntu-latest @@ -83,6 +112,7 @@ jobs: exit 1 fi BundleMetadataContract: + needs: ReleaseLock name: Validate bundle metadata contract runs-on: ubuntu-latest steps: @@ -94,14 +124,15 @@ jobs: - name: Install uv uses: astral-sh/setup-uv@v5 - name: Install package - run: uv pip install --system . + run: uv sync --locked --extra dev - name: Install bundle validation tooling # Pin the test-only bundle contract dependency until policyengine-bundles # has published releases suitable for ordinary dependency specifiers. - run: uv pip install --system "policyengine-bundles @ git+https://github.com/PolicyEngine/policyengine-bundles@8ae9f56fefcf89f69b8a7e3bc49928509c6207be" + run: uv pip install --python .venv/bin/python "policyengine-bundles @ git+https://github.com/PolicyEngine/policyengine-bundles@8ae9f56fefcf89f69b8a7e3bc49928509c6207be" - name: Validate runtime metadata contract - run: python -m pytest policyengine_us/tests/test_build_metadata.py + run: uv run --no-sync python -m pytest policyengine_us/tests/test_build_metadata.py Python-Compat: + needs: ReleaseLock name: Install + smoke-import (py${{ matrix.python-version }}) runs-on: ubuntu-latest strategy: @@ -117,14 +148,15 @@ jobs: - name: Install uv uses: astral-sh/setup-uv@v8.1.0 - name: Install package - run: uv pip install --system . pytest + run: uv sync --locked --python "${{ matrix.python-version }}" --extra dev - name: Smoke-import run: | - python -c "from policyengine_us import CountryTaxBenefitSystem, Simulation, Microsimulation; print('import OK')" + uv run --no-sync --python "${{ matrix.python-version }}" python -c "from policyengine_us import CountryTaxBenefitSystem, Simulation, Microsimulation; print('import OK')" - name: Dataset copy/extension tests (real pandas semantics per Python version) # Fixture-based and requires no dataset downloads. - run: python -m pytest policyengine_us/tests/microsimulation/data/ -q + run: uv run --no-sync --python "${{ matrix.python-version }}" python -m pytest policyengine_us/tests/microsimulation/data/ -q Quick-Feedback: + needs: ReleaseLock name: Quick Feedback (Selective Tests + Coverage) runs-on: ubuntu-latest env: @@ -150,7 +182,7 @@ jobs: - name: Install uv uses: astral-sh/setup-uv@v8.1.0 - name: Install dependencies - run: uv sync --extra dev + run: uv sync --locked --extra dev - name: Turn off default branching shell: bash run: ./update_itemization.sh @@ -161,7 +193,7 @@ jobs: GITHUB_BASE_REF: ${{ github.base_ref }} run: | echo "Running selective tests based on changed files..." - uv run python policyengine_us/tests/run_selective_tests.py --base-branch ${{ github.base_ref }} --coverage + uv run --no-sync python policyengine_us/tests/run_selective_tests.py --base-branch ${{ github.base_ref }} --coverage - name: Generate coverage report id: coverage if: always() @@ -169,10 +201,10 @@ jobs: echo "📊 Coverage Report" echo "==================" # Show files below 100% coverage (won't fail CI) - uv run coverage report --fail-under=0 --skip-covered --skip-empty || echo "No coverage data available." + uv run --no-sync coverage report --fail-under=0 --skip-covered --skip-empty || echo "No coverage data available." # Count files below 100% - TOTAL_FILES=$(uv run coverage report --skip-covered --skip-empty 2>/dev/null | grep -E "^[a-zA-Z].*\.py " | wc -l || echo "0") + TOTAL_FILES=$(uv run --no-sync coverage report --skip-covered --skip-empty 2>/dev/null | grep -E "^[a-zA-Z].*\.py " | wc -l || echo "0") if [ "$TOTAL_FILES" -gt "0" ]; then echo "" @@ -184,7 +216,7 @@ jobs: fi # Generate XML for Codecov - uv run coverage xml || echo "No coverage data to generate XML report." + uv run --no-sync coverage xml || echo "No coverage data to generate XML report." # Expose whether we have coverage for later steps if [ -f coverage.xml ]; then echo "has_coverage=true" >> "$GITHUB_OUTPUT" @@ -204,6 +236,7 @@ jobs: - name: Build package run: uv build Baseline: + needs: ReleaseLock name: Full Suite - Baseline (${{ matrix.group }}) runs-on: ubuntu-latest timeout-minutes: 60 @@ -248,13 +281,14 @@ jobs: - name: Install uv uses: astral-sh/setup-uv@v8.1.0 - name: Install dependencies - run: uv sync --extra dev + run: uv sync --locked --extra dev - name: Turn off default branching shell: bash run: bash ./update_itemization.sh - name: Run Baseline YAML tests (${{ matrix.group }}) - run: uv run ${{ matrix.cmd }} + run: uv run --no-sync ${{ matrix.cmd }} HouseholdAPIPartners: + needs: ReleaseLock name: Household API Partners runs-on: ubuntu-latest timeout-minutes: 60 @@ -271,13 +305,14 @@ jobs: - name: Install uv uses: astral-sh/setup-uv@v8.1.0 - name: Install dependencies - run: uv sync --extra dev + run: uv sync --locked --extra dev - name: Turn off default branching shell: bash run: bash ./update_itemization.sh - name: Run Household API Partners YAML tests - run: uv run make test-yaml-no-structural-other-partners + run: uv run --no-sync make test-yaml-no-structural-other-partners Contrib: + needs: ReleaseLock name: Full Suite - Contrib (${{ matrix.group }}) runs-on: ubuntu-latest timeout-minutes: 60 @@ -316,16 +351,17 @@ jobs: - name: Install uv uses: astral-sh/setup-uv@v8.1.0 - name: Install dependencies - run: uv sync --extra dev + run: uv sync --locked --extra dev - name: Turn off default branching shell: bash run: bash ./update_itemization.sh - name: Run Contrib YAML tests (${{ matrix.group }}) - run: uv run make ${{ matrix.target }} + run: uv run --no-sync make ${{ matrix.target }} - name: Run Contrib Python tests if: ${{ matrix.group == 'other-shard-1' }} - run: uv run make test-policy-contrib-python + run: uv run --no-sync make test-policy-contrib-python Rest: + needs: ReleaseLock name: Full Suite - Rest (Python + variables) runs-on: ubuntu-latest timeout-minutes: 60 @@ -342,16 +378,17 @@ jobs: - name: Install uv uses: astral-sh/setup-uv@v8.1.0 - name: Install dependencies - run: uv sync --extra dev + run: uv sync --locked --extra dev - name: Turn off default branching shell: bash run: bash ./update_itemization.sh - name: Run Python tests (code_health, core, utilities, top-level) - run: uv run make test-other-python + run: uv run --no-sync make test-other-python - name: Run tests/variables YAML tests if: always() - run: uv run make test-yaml-variables + run: uv run --no-sync make test-yaml-variables Microsimulation: + needs: ReleaseLock name: Full Suite - Microsimulation runs-on: ubuntu-latest timeout-minutes: 60 @@ -368,7 +405,7 @@ jobs: - name: Install uv uses: astral-sh/setup-uv@v8.1.0 - name: Install dependencies - run: uv sync --extra dev + run: uv sync --locked --extra dev - name: Turn off default branching shell: bash run: bash ./update_itemization.sh @@ -376,4 +413,4 @@ jobs: # was OOM-killed at ~94% on run 28698452678 once microsimulation # peaks stacked on the accumulated python-test RSS. - name: Run microsimulation tests - run: uv run make test-microsimulation + run: uv run --no-sync make test-microsimulation diff --git a/.github/workflows/push.yaml b/.github/workflows/push.yaml index 1a5363d47f9..ab6ab6af2b9 100644 --- a/.github/workflows/push.yaml +++ b/.github/workflows/push.yaml @@ -4,7 +4,39 @@ on: branches: [main] jobs: + ReleaseLock: + name: Check committed registry lock + if: | + (github.repository == 'PolicyEngine/policyengine-us') + && (github.event.head_commit.message == 'Update PolicyEngine US') + runs-on: ubuntu-latest + steps: + - uses: actions/checkout@v6 + - name: Set up Python + uses: actions/setup-python@v6 + with: + python-version: "3.14" + - name: Install uv + uses: astral-sh/setup-uv@v8.1.0 + - name: Check committed registry lock + run: python .github/release_lock.py --committed + ReleaseLockTests: + name: Release lock guard tests + runs-on: ubuntu-latest + steps: + - uses: actions/checkout@v6 + - name: Set up Python + uses: actions/setup-python@v6 + with: + python-version: "3.14" + - name: Install uv + uses: astral-sh/setup-uv@v8.1.0 + - name: Test guard without importing the country model + env: + RELEASE_LOCK_REAL_UV: "1" + run: python -m unittest discover -s .github/tests -p test_release_lock.py -v Lint: + needs: ReleaseLock runs-on: ubuntu-latest if: | (github.repository == 'PolicyEngine/policyengine-us') @@ -14,10 +46,11 @@ jobs: - name: Install uv uses: astral-sh/setup-uv@v8.1.0 - name: Install dependencies - run: uv sync --extra dev + run: uv sync --locked --extra dev - name: Check formatting - run: uv run ruff format --check . + run: uv run --no-sync ruff format --check . versioning: + needs: ReleaseLockTests name: Update versioning if: | (github.repository == 'PolicyEngine/policyengine-us') @@ -40,23 +73,50 @@ jobs: uses: actions/setup-python@v6 with: python-version: 3.14 - - name: Install dependencies - run: uv pip install towncrier --system - - name: Bump version and build changelog + - name: Check committed registry lock before bump + run: python .github/release_lock.py --committed + - name: Bump version once + run: python .github/bump_version.py + - name: Build changelog with isolated registry tooling run: | - python .github/bump_version.py - VERSION=$(python -c "import re; print(re.search(r'version = \"(.+?)\"', open('pyproject.toml').read()).group(1))") - towncrier build --yes --version "$VERSION" + python - <<'PYTHON' + import os + import subprocess + import sys + import tomllib + from pathlib import Path + + version = tomllib.loads(Path("pyproject.toml").read_text())["project"]["version"] + env = { + key: value + for key, value in os.environ.items() + if not key.startswith(("UV_", "PIP_")) + } + subprocess.run( + [ + "uv", "--no-config", "tool", "run", "--isolated", + "--no-env-file", "--no-sources", + "--default-index", "https://pypi.org/simple", + "--python", sys.executable, + "towncrier", "build", "--yes", "--version", version, + ], + env=env, + check=True, + ) + PYTHON + - name: Refresh and check release version in registry lock + run: python .github/release_lock.py --refresh - name: Update changelog uses: EndBug/add-and-commit@v10 with: - add: "." + add: "pyproject.toml uv.lock CHANGELOG.md changelog.d" committer_name: Github Actions[bot] author_name: Github Actions[bot] message: Update PolicyEngine US github_token: ${{ steps.app-token.outputs.token }} fetch: false Baseline: + needs: ReleaseLock name: Full Suite - Baseline (${{ matrix.group }}) runs-on: ubuntu-latest timeout-minutes: 60 @@ -104,13 +164,14 @@ jobs: - name: Install uv uses: astral-sh/setup-uv@v8.1.0 - name: Install dependencies - run: uv sync --extra dev + run: uv sync --locked --extra dev - name: Turn off default branching shell: bash run: bash ./update_itemization.sh - name: Run Baseline YAML tests (${{ matrix.group }}) - run: uv run ${{ matrix.cmd }} + run: uv run --no-sync ${{ matrix.cmd }} HouseholdAPIPartners: + needs: ReleaseLock name: Household API Partners runs-on: ubuntu-latest timeout-minutes: 60 @@ -130,13 +191,14 @@ jobs: - name: Install uv uses: astral-sh/setup-uv@v8.1.0 - name: Install dependencies - run: uv sync --extra dev + run: uv sync --locked --extra dev - name: Turn off default branching shell: bash run: bash ./update_itemization.sh - name: Run Household API Partners YAML tests - run: uv run make test-yaml-no-structural-other-partners + run: uv run --no-sync make test-yaml-no-structural-other-partners Contrib: + needs: ReleaseLock name: Full Suite - Contrib (${{ matrix.group }}) runs-on: ubuntu-latest timeout-minutes: 60 @@ -178,16 +240,17 @@ jobs: - name: Install uv uses: astral-sh/setup-uv@v8.1.0 - name: Install dependencies - run: uv sync --extra dev + run: uv sync --locked --extra dev - name: Turn off default branching shell: bash run: bash ./update_itemization.sh - name: Run Contrib YAML tests (${{ matrix.group }}) - run: uv run make ${{ matrix.target }} + run: uv run --no-sync make ${{ matrix.target }} - name: Run Contrib Python tests if: ${{ matrix.group == 'other-shard-1' }} - run: uv run make test-policy-contrib-python + run: uv run --no-sync make test-policy-contrib-python Rest: + needs: ReleaseLock name: Full Suite - Rest (Python + variables) runs-on: ubuntu-latest timeout-minutes: 60 @@ -207,16 +270,17 @@ jobs: - name: Install uv uses: astral-sh/setup-uv@v8.1.0 - name: Install dependencies - run: uv sync --extra dev + run: uv sync --locked --extra dev - name: Turn off default branching shell: bash run: bash ./update_itemization.sh - name: Run Python tests (code_health, core, utilities, top-level) - run: uv run make test-other-python + run: uv run --no-sync make test-other-python - name: Run tests/variables YAML tests if: always() - run: uv run make test-yaml-variables + run: uv run --no-sync make test-yaml-variables Microsimulation: + needs: ReleaseLock name: Full Suite - Microsimulation runs-on: ubuntu-latest timeout-minutes: 60 @@ -236,7 +300,7 @@ jobs: - name: Install uv uses: astral-sh/setup-uv@v8.1.0 - name: Install dependencies - run: uv sync --extra dev + run: uv sync --locked --extra dev - name: Turn off default branching shell: bash run: bash ./update_itemization.sh @@ -244,13 +308,13 @@ jobs: # was OOM-killed at ~94% on run 28698452678 once microsimulation # peaks stacked on the accumulated python-test RSS. - name: Run microsimulation tests - run: uv run make test-microsimulation + run: uv run --no-sync make test-microsimulation Publish: runs-on: ubuntu-latest if: | (github.repository == 'PolicyEngine/policyengine-us') && (github.event.head_commit.message == 'Update PolicyEngine US') - needs: [Baseline, HouseholdAPIPartners, Contrib, Rest, Microsimulation] + needs: [ReleaseLock, ReleaseLockTests, Baseline, HouseholdAPIPartners, Contrib, Rest, Microsimulation] steps: - name: Checkout repo uses: actions/checkout@v6 @@ -260,10 +324,12 @@ jobs: python-version: 3.14 - name: Install uv uses: astral-sh/setup-uv@v8.1.0 + - name: Check committed registry lock before publication + run: python .github/release_lock.py --committed - name: Install package - run: uv pip install -e ".[dev]" --system + run: uv sync --locked --extra dev - name: Build package - run: make + run: uv run --no-sync make - name: Publish a Python distribution to PyPI uses: pypa/gh-action-pypi-publish@release/v1 with: diff --git a/Makefile b/Makefile index 738bb9be21b..ddf009d290b 100644 --- a/Makefile +++ b/Makefile @@ -1,7 +1,8 @@ # Shorthand for the batch runner so each target is easier to read. -# --mode per-subdir = each immediate subdir runs in its own subprocess -# (loose yamls get a trailing batch). New subdirs -# auto-route, no Makefile edit needed. +# --mode per-subdir = each immediate subdir runs in its own subprocess, +# split by reform-combo weight when it exceeds the +# batcher's budget (loose yamls get a trailing batch). +# New subdirs auto-route, no Makefile edit needed. # --mode per-file = each yaml runs in its own subprocess. Used for # microsim-heavy folders where one file per subprocess # is needed to keep peak RAM under the 16 GB runner. @@ -70,7 +71,8 @@ test-yaml-structural-other-shard-3: # one peak resident at a time. $(BATCH) $(TESTS)/policy/contrib/refundable_credit_conversion --mode per-file --workers 1 test-yaml-structural-congress: - # One subprocess per congress proposal; new proposals auto-route. + # One subprocess per congress proposal, split when a proposal's reform + # combos exceed the batcher's budget; new proposals auto-route. # --workers 1: congress OOM'd two-wide on CI run 28698452678 — the # romney batch alone peaked 7.1 GB. $(BATCH) $(TESTS)/policy/contrib/congress --mode per-subdir --workers 1 diff --git a/changelog.d/country-release-lock.fixed.md b/changelog.d/country-release-lock.fixed.md new file mode 100644 index 00000000000..9e20681a66b --- /dev/null +++ b/changelog.d/country-release-lock.fixed.md @@ -0,0 +1 @@ +Prevent automatic releases from committing stale or non-PyPI dependency locks, and reject dependency drift when refreshing the package version. diff --git a/changelog.d/spm-canonical-default.breaking.md b/changelog.d/spm-canonical-default.breaking.md new file mode 100644 index 00000000000..7abe3673dc2 --- /dev/null +++ b/changelog.d/spm-canonical-default.breaking.md @@ -0,0 +1 @@ +Require county FIPS by default, or an explicit national or fixed SPM area selection, for SPM measurement, and for resource calculations (household net income, benefits, marginal tax rates) only where a unit's housing assistance is positive: the housing cap consults the canonical housing portion for assisted units alone, so partner and state-program outputs for the other units never depend on SPM geography. Replace country threshold extrapolation with spm-calculator 1.0.0's canonical 2022–2035 amounts; unavailable years fail. Population datasets must supply observed county inputs and source-backed SPM independence roles instead of stored formula-owned SPM outputs, which the loader now rejects; the legacy policyengine-us-data CPS files supply none of these, so resource and poverty outputs are no longer available over them. Preserve user reforms, isolate simulation receipts and bind cloned and baseline holders to their own policy variables. diff --git a/docs/spm.md b/docs/spm.md new file mode 100644 index 00000000000..231dc38b197 --- /dev/null +++ b/docs/spm.md @@ -0,0 +1,78 @@ +# Supplemental Poverty Measure + +The country model uses the installed `spm-calculator` artifact for all five SPM +measurement amounts. It does not extrapolate thresholds with country CPI or income +parameters. Each supported year and scenario comes from the verified artifact. +The installed 1.0.0 artifact covers 2022 through 2035; unavailable years fail. +Taxes, benefits, resources and the housing-assistance cap +remain country-model formulas. The cap uses the canonical housing amount before +the final model storage conversion. + +`Simulation`, `Microsimulation` and `CountryTaxBenefitSystem` accept a serializable +`spm` mapping. Its fields are `forecast_content_sha256`, `scenario`, +`geography_kind`, `geography_id`, `county_vintage` and `as_of`. The hash, when given, +must match the installed artifact. The default scenario is the artifact's declared +default; the default county vintage is `"2020"`. + +County geography is the default for both household and population calculations. +A local household must supply a valid `county_fips`. State alone does not identify +an SPM area. Missing counties raise `SPM_GEOGRAPHY_REQUIRED`; invalid or unavailable +counties raise `SPM_GEOGRAPHY_UNAVAILABLE`. There is no first-county, congressional +district or national fallback in SPM measurement. + +These geography errors occur only when calculating an SPM measurement or a +dependent resource, such as the housing-assistance cap for units with housing +assistance; units with none are capped at zero without consulting the +measurement. A state-only tax request can still run. SPM reads the input-only `county_fips` variable and ignores any +county inferred or cached by other tax or benefit formulas. Geography and +composition errors are `SPMInputError` instances with `code` and `to_dict()`. + +A caller can consciously select national measurement: + +```python +simulation = Simulation(situation=household, spm={"geography_kind": "national"}) +``` + +Or select one artifact area with `{"geography_kind": "metro", "geography_id": +""}`. A public application should display the national/local choice and +pass it explicitly; it should return structured validation errors when requested +local inputs are missing. This selection affects SPM measurement, not the +geography inputs used by other benefit programs. + +`geography_id` applies only to a fixed metro selection; county mode uses each +household's observed `county_fips`. `as_of` is an ISO calendar date constraining +the artifact's information date. Configuration accepts no external forecast path +or consumer-side extrapolation policy. + +`simulation.spm_config` returns fully resolved JSON-compatible settings, including +the artifact content hash. `simulation.spm_provenance()` returns detached runtime, +year, scenario and geography receipts. Each simulation has private receipt state; +policy-reform baselines retain the same measurement selection. Starting a new +simulation from a supplied system preserves its reform variables and starts fresh +receipts. Cloning a calculated simulation retains receipts for cached values and +detaches future receipts. Branches sharing policy parameters also have separate +SPM receipt state. + +SPM adult counts use supplied age and a source-backed independence role. Explicit +household head/spouse primitives support household scenarios; the model never +guesses those roles from age ordering. A person counts as an SPM adult at age 18, +or at age 15 or above with `is_spm_independent_minor_role`. That role defaults to +the input-only `is_household_head | is_household_spouse` and can be supplied from +source data. A unit with no classified adult raises +`SPM_COMPOSITION_REQUIRED`. Generic age-based adult/child counts and benefit +eligibility are unchanged. The dataset loader rejects stored formula-owned SPM +outputs; observed source results should use report-only column names. It does not +delete or silently recalculate over supplied derived inputs. + +`Microsimulation()` defaults to the immutable dataset URI +`hf://datasets/policyengine/populace-us/populace_us_2024.h5@populace-us-2024-spm-20260909`. +This dataset supplies observed county inputs and source-backed independence roles. +The canonical country release must wait until that exact tag and its certified +bytes exist and pass independent readback. The unpublished candidate embeds the +same URI; local candidate checks do not establish production default availability. + +The published country wheel requires exactly `spm-calculator==1.0.0`; it contains +no Git or local-path dependency. For coordinated development before registry +resolution is available, install the local calculator wheel into an isolated +environment and run tests with `uv run --no-sync`. Refresh the registry lockfile +when the calculator release is available to the resolver. diff --git a/docs/usage/microsimulation.md b/docs/usage/microsimulation.md index 34fcfa311bc..95db177874a 100644 --- a/docs/usage/microsimulation.md +++ b/docs/usage/microsimulation.md @@ -91,18 +91,37 @@ hh_eitc = sim.calc("eitc", map_to="household") ## Available datasets -PolicyEngine US provides datasets hosted on HuggingFace at `hf://policyengine/policyengine-us-data/`. +### The default build -### National datasets +The default is the certified Populace build, pinned by build id and hosted as a +HuggingFace *dataset* repository: ```python -# Default: Enhanced CPS 2024 (includes imputed wealth and other enhancements) +# Default: the certified Populace build named in DEFAULT_DATASET. sim = Microsimulation() -# Equivalent to: -sim = Microsimulation(dataset="hf://policyengine/policyengine-us-data/enhanced_cps_2024.h5") +``` + +### Supplying another population + +A population file has to satisfy the SPM input contract: it supplies primitive +inputs, including observed `county_fips` codes as five-digit strings and +source-backed `is_spm_independent_minor_role` values, and it must not store +formula-owned SPM outputs such as `spm_unit_spm_threshold`. Any observed Census +measurement is retained under a separate report-only name. + +The legacy files under `hf://policyengine/policyengine-us-data/` predate that +contract, so resource and poverty outputs are not available over them: -# Basic CPS 2023 -sim = Microsimulation(dataset="hf://policyengine/policyengine-us-data/cps_2023.h5") +- `cps_2023.h5` stores `spm_unit_spm_threshold`, so the loader rejects it. +- `enhanced_cps_2024.h5` loads and computes tax variables, but carries no SPM + independence roles, so SPM units consisting of one 15-to-17-year-old classify + no measurement adult and every resource output over the file fails closed. + +A household simulation that has no county input can select an SPM area +explicitly instead: + +```python +sim = Microsimulation(dataset=..., spm={"geography_kind": "national"}) ``` ### Filtering by geography diff --git a/policyengine_us/data/economic_assumptions.py b/policyengine_us/data/economic_assumptions.py index b57284c253a..5d9c1a44da4 100644 --- a/policyengine_us/data/economic_assumptions.py +++ b/policyengine_us/data/economic_assumptions.py @@ -37,7 +37,6 @@ "social_security": "calibration.gov.irs.soi.social_security", "s_corp_income": "calibration.gov.irs.soi.partnership_s_corp_income", "spm_unit_weight": "calibration.gov.census.populations.total", - "spm_unit_spm_threshold": DEFAULT_MICRODATA_UPRATING, "state_and_local_sales_or_income_tax": DEFAULT_MICRODATA_UPRATING, "sstb_self_employment_income": "calibration.gov.irs.soi.self_employment_income", "taxable_pension_income": "calibration.gov.irs.soi.taxable_pension_income", diff --git a/policyengine_us/spm.py b/policyengine_us/spm.py new file mode 100644 index 00000000000..bedc086632e --- /dev/null +++ b/policyengine_us/spm.py @@ -0,0 +1,356 @@ +"""Explicit, serializable selection of the canonical SPM forecast. + +The calculator owns measurement formulas. The country owns model construction, +input handling and policy resources. Geography is never inferred from a missing +county: callers must explicitly select national or a particular metropolitan area. +""" + +import re +from collections.abc import Mapping +from copy import copy, deepcopy +from functools import lru_cache +from inspect import signature + +from policyengine_core.simulations import Simulation as CoreSimulation +from policyengine_core.taxbenefitsystems import TaxBenefitSystem +from spm_calculator.errors import SPMInputError +from spm_calculator.policyengine_adapter import ( + FORMULA_OWNED_INPUTS, + PolicyEngineSPMProvider, +) +from spm_calculator.rolling_forecast import load_forecast + + +CONFIG_FIELDS = frozenset( + { + "forecast_content_sha256", + "scenario", + "geography_kind", + "geography_id", + "county_vintage", + "as_of", + } +) + +COUNTY_FIPS_PATTERN = re.compile(r"[0-9]{5}") + +COUNTY_INPUT_FIX = ( + 'send county_fips as a five-digit string (for example "06037"), or select ' + 'geography_kind="national" in the spm configuration' +) + + +def is_county_fips(value): + """Accept only a five-digit county FIPS code. + + Integers, pandas missing values and truncated codes all arrive here as + strings, because ``county_fips`` is a string variable and the model casts + every input to its own dtype. A within-state CPS code such as ``5`` and a + missing value such as ``nan`` are not county FIPS codes, so they must fail + as an absent county rather than as an unrecognised one. + """ + if value is None: + return False + if isinstance(value, bytes): + value = value.decode() + return COUNTY_FIPS_PATTERN.fullmatch(str(value)) is not None + + +class CountyRequiringSPMProvider(PolicyEngineSPMProvider): + """Name the caller's fix when a county selection has no usable county. + + The calculator reports an absent county only for ``None`` and the empty + string, and reports anything else as an unavailable county. A legacy + population file storing the CPS within-state integer code, or a household + request sending an integer or a missing value, would otherwise be told its + county assignment is unavailable, which points at the artifact instead of + at the input. + """ + + def calculate_unit(self, *, year, adults, children, tenure, county_fips=None): + if self.geography_kind == "county" and not is_county_fips(county_fips): + raise SPMInputError( + "SPM_GEOGRAPHY_REQUIRED", + f"County selection has no county FIPS input ({county_fips!r}): " + f"{COUNTY_INPUT_FIX}", + ) + return super().calculate_unit( + year=year, + adults=adults, + children=children, + tenure=tenure, + county_fips=county_fips, + ) + + +def masked_policyengine_amount(unit, period, field, mask): + """Return canonical float64 amounts for the units selected by ``mask``. + + The calculator's ``policyengine_amount`` evaluates every unit in the + population, so a resource formula that only needs the SPM housing portion + for units with housing assistance to cap would otherwise demand SPM + geography and composition from units that never use them. Evaluate the + selected units only, through the same provider path, so the county and + composition requirements, the typed errors and the receipts attach to + exactly the units whose result depends on the measurement. Unselected + units return 0.0 and record nothing. + """ + import numpy as np + + fields = ( + "reference_threshold", + "unadjusted_threshold", + "geographic_factor", + "threshold", + "housing_portion", + ) + if field not in fields: + raise ValueError(f"Unknown SPM amount field: {field}") + index = fields.index(field) + mask = np.asarray(mask, dtype=bool) + result = np.zeros(mask.shape, dtype=np.float64) + if not mask.any(): + return result + bound = unit.simulation.tax_benefit_system.spm_forecast_provider + adults = np.asarray(unit("spm_measurement_adults", period))[mask] + children = np.asarray(unit("spm_measurement_children", period))[mask] + if np.any(adults < 1): + raise SPMInputError( + "SPM_COMPOSITION_REQUIRED", + "SPM unit has no classified adult: supply source-backed independence or household head/spouse structure", + ) + tenures = np.asarray(unit("spm_unit_tenure_type", period).decode_to_str())[mask] + if bound.geography_kind == "county": + counties = np.asarray(unit.household("county_fips", period))[mask] + else: + counties = [None] * len(adults) + rows = [ + bound._amounts( + int(period.start.year), + int(a), + int(k), + str(t).lower(), + None if c is None else (c.decode() if isinstance(c, bytes) else str(c)), + ) + for a, k, t, c in zip(adults, children, tenures, counties) + ] + result[mask] = [row[index] for row in rows] + return result + + +@lru_cache(maxsize=1) +def _installed_forecast(): + """Verify the bundled immutable artifact once; never download data.""" + return load_forecast() + + +def create_spm_provider(config=None): + """Resolve a public configuration against this installed calculator.""" + if config is None: + config = {} + if not isinstance(config, Mapping): + raise TypeError("spm must be a mapping of explicit forecast settings") + unknown = set(config) - CONFIG_FIELDS + if unknown: + raise ValueError(f"Unknown spm settings: {', '.join(sorted(unknown))}") + forecast = _installed_forecast() + expected = config.get("forecast_content_sha256") + if expected is not None and expected != forecast.content_sha256: + raise ValueError( + "The installed SPM forecast does not match forecast_content_sha256" + ) + return CountyRequiringSPMProvider( + forecast=forecast, + **{ + key: value + for key, value in config.items() + if key != "forecast_content_sha256" + }, + ) + + +def spm_config(provider): + """Return detached JSON-compatible settings suitable for a bundle receipt.""" + return { + "forecast_content_sha256": provider.forecast.content_sha256, + "scenario": provider.scenario, + "geography_kind": provider.geography_kind, + "geography_id": provider.geography_id, + "county_vintage": provider.county_vintage, + "as_of": provider.as_of, + } + + +def share_spm_policy(system): + """Isolate receipts and variable registration without rebuilding policy. + + An ordinary simulation applies no user reform, so it needs private receipts + and a private variable registry, not a private copy of the policy itself. + It is also new rather than cloned, so no previous receipt belongs to it. + Core's TaxBenefitSystem.clone() rebuilds the whole parameter tree node by + node and empties both at-instant caches, and this country then deep-copies + every variable object on top; doing that per household simulation throws + away the shared instance's warm parameter caches and lands on household API + request latency. + + Share the parameter tree and its at-instant caches, and reuse the variable + objects. Every core operation that a reform performs on a variable + (add_variable, replace_variable, update_variable, neutralize_variable, + annualize_variable) rebinds ``variables[name]`` to a newly constructed + object rather than mutating the registered one, so a private dict is enough + to keep this simulation's registration - including the structural reform + re-applied at its own start instant - out of the shared instance. + ``test_ordinary_simulation_shares_default_policy_state`` enforces that + invariant against the shared instance itself. + """ + policy = copy(system) + policy.variables = dict(system.variables) + policy.spm_forecast_provider = system.spm_forecast_provider.snapshot( + copy_receipts=False + ) + return policy + + +def clone_spm_system(system, *, copy_receipts=True): + """Clone policy state without reconstructing partially defined reforms. + + Core's Variable.clone() calls the variable class with no baseline metadata, + losing instance changes such as neutralization and inherited reform fields. + Let core clone parameters/entities, then copy the actual variable state. + """ + policy = copy(system) + policy.variables = {} + cloned = TaxBenefitSystem.clone(policy) + by_key = {entity.key: entity for entity in cloned.entities} + memo = {id(system): cloned, id(system.parameters): cloned.parameters} + for entity in system.entities: + memo[id(entity)] = by_key[entity.key] + for variable in system.variables.values(): + memo[id(variable.entity)] = by_key[variable.entity.key] + cloned.variables = deepcopy(system.variables, memo) + cloned.spm_forecast_provider = system.spm_forecast_provider.snapshot( + copy_receipts=copy_receipts + ) + return cloned + + +class SPMSimulationMixin: + """Keep forecast selection and receipts private to each simulation.""" + + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + # Core switches the baseline system after cloning its populations. + self._rebind_holders() + + def _rebind_holders(self): + """Bind populations and cached holders to their branch's policy state.""" + self.tax_benefit_system.simulation = self + variables = self.tax_benefit_system.variables + for entity in self.tax_benefit_system.entities: + population = self.populations[entity.key] + population.entity = entity + for name, holder in list(population._holders.items()): + if name in variables: + holder.variable = variables[name] + else: + # A reform-only variable has no holder in baseline policy. + del population._holders[name] + # Core shallow-copies this metadata; detach it before pruning inputs + # that exist only in the reform, so later cache invalidation is valid. + self._user_input_keys = { + key for key in self._user_input_keys if key[0] in variables + } + self.input_variables = [ + name for name in self.input_variables if name in variables + ] + for branch in self.branches.values(): + branch._rebind_holders() + + def apply_reform(self, reform): + super().apply_reform(reform) + self._rebind_holders() + + @property + def spm_config(self): + return spm_config(self.tax_benefit_system.spm_forecast_provider) + + def spm_provenance(self): + return self.tax_benefit_system.spm_forecast_provider.provenance() + + def _prepare_spm_system(self, args, kwargs, config, start_instant): + # Register the shared calculator variables before the input builder sees + # new primitive role inputs. Cloning also isolates calculation receipts. + arguments = signature(CoreSimulation.__init__).bind_partial( + self, *args, **kwargs + ) + supplied = arguments.arguments.get("tax_benefit_system") + reform = arguments.arguments.get("reform") + if supplied is None: + if reform is not None: + chosen = self.default_tax_benefit_system( + reform=reform, spm=config, start_instant=start_instant + ) + else: + # No reform and no supplied system: nothing distinguishes this + # simulation's policy from the shared instance's, so keep its + # warm parameter caches instead of rebuilding them. + chosen = share_spm_policy(self.default_tax_benefit_system_instance) + elif reform is not None: + # Core applies the reform set to whatever system it is handed, so + # the caller's own system must not be the one it reforms. + chosen = clone_spm_system(supplied, copy_receipts=False) + else: + # A caller that builds one system and runs many households through + # it - the household API among them - keeps that system's warm + # caches; only receipts and variable registration are private. + chosen = share_spm_policy(supplied) + # This is a new simulation, unlike clone() of an already calculated + # simulation, so no previous calculation receipt belongs to it. + if config is not None: + chosen.spm_forecast_provider = create_spm_provider(config) + arguments.arguments["tax_benefit_system"] = chosen + + # Core constructs a baseline branch for policy reforms. It must use the + # same SPM selection while retaining baseline tax/benefit policy. + if reform is not None: + # The baseline branch holds unreformed policy, which is the shared + # instance's own policy, so it shares rather than rebuilds it. + baseline = share_spm_policy(self.default_tax_benefit_system_instance) + baseline.spm_forecast_provider = chosen.spm_forecast_provider.snapshot() + self.default_tax_benefit_system_instance = baseline + # Country dataset interception also needs to see positional datasets. + return (), { + key: value for key, value in arguments.arguments.items() if key != "self" + } + + def clone(self, debug=False, trace=False, clone_tax_benefit_system=True): + """Retain cached-result receipts while isolating future calculations.""" + cloned = super().clone(debug=debug, trace=trace, clone_tax_benefit_system=False) + if clone_tax_benefit_system: + cloned.tax_benefit_system = clone_spm_system(self.tax_benefit_system) + else: + # Core branches may share policy state, but each calculation's + # receipt belongs to that simulation's provider. + cloned.tax_benefit_system = copy(self.tax_benefit_system) + cloned.tax_benefit_system.spm_forecast_provider = ( + self.tax_benefit_system.spm_forecast_provider.snapshot( + copy_receipts=True + ) + ) + cloned._rebind_holders() + return cloned + + def set_input(self, variable_name, period, value): + # Core's loader calls set_input for every dataset format. Reject saved + # measurement/resource outputs here without loading the population twice. + # Household situation overrides remain available for model unit tests; + # public consumers validate their household input contract separately. + if ( + getattr(self, "is_over_dataset", False) + and variable_name in FORMULA_OWNED_INPUTS + ): + raise ValueError( + f"Dataset supplies formula-owned SPM output {variable_name}. " + "Use primitive inputs and retain observed outputs under report-only names." + ) + return super().set_input(variable_name, period, value) diff --git a/policyengine_us/system.py b/policyengine_us/system.py index 0c0060ed8e3..6b83b8b3acd 100644 --- a/policyengine_us/system.py +++ b/policyengine_us/system.py @@ -37,6 +37,12 @@ ) from typing import Annotated +from spm_calculator.policyengine_adapter import build_policyengine_variables +from policyengine_us.spm import ( + SPMSimulationMixin, + clone_spm_system, + create_spm_provider, +) COUNTRY_DIR = Path(__file__).parent @@ -46,7 +52,7 @@ # Certified Populace build (primary-source US microdata), pinned by build id. # Populace ships from a Hugging Face *dataset* repo, hence the `hf://datasets/` # prefix handled in `_resolve_dataset_path`. -DEFAULT_DATASET = "hf://datasets/policyengine/populace-us/populace_us_2024.h5@populace-us-2024-c86a631-6e1bcd0271a5-20260619T002242Z" +DEFAULT_DATASET = "hf://datasets/policyengine/populace-us/populace_us_2024.h5@populace-us-2024-spm-20260909" class CountryTaxBenefitSystem(TaxBenefitSystem): @@ -82,8 +88,11 @@ def __init__( start_instant: Annotated[ str, "ISO date format YYYY-MM-DD" ] = DEFAULT_START_DATE, + spm: dict | None = None, ): super().__init__(entities, reform=reform) + self.spm_forecast_provider = create_spm_provider(spm) + self.add_variables(*build_policyengine_variables()) self.load_parameters(COUNTRY_DIR / "parameters") self.add_abolition_parameters() self.parameters = set_irs_uprating_parameter(self.parameters) @@ -125,6 +134,9 @@ def __init__( self.add_variables(*create_50_state_variables()) + def clone(self): + return clone_spm_system(self) + system = CountryTaxBenefitSystem() @@ -158,7 +170,7 @@ def _backfill_state_code_from_str(simulation): state_code_str.delete_arrays(known_period) -class Simulation(CoreSimulation): +class Simulation(SPMSimulationMixin, CoreSimulation): """ A simulation of the tax-benefit system for the United States, defined against the base simulation class in the -core package. @@ -184,6 +196,9 @@ def __init__(self, *args, **kwargs): start_instant: Annotated[str, "ISO date format YYYY-MM-DD"] = kwargs.pop( "start_instant", DEFAULT_START_DATE ) + args, kwargs = self._prepare_spm_system( + args, kwargs, kwargs.pop("spm", None), start_instant + ) super().__init__(*args, **kwargs) reform = create_structural_reforms_from_parameters( @@ -234,6 +249,25 @@ def __init__(self, *args, **kwargs): _backfill_state_code_from_str(self) +def _download_or_explain(dataset_str, download): + """Name the unresolved dataset URI instead of re-raising a Hub exception. + + A build id that is not published yet, a renamed repository and an offline + cache miss all surface from ``huggingface_hub`` as transport-level errors + that never mention which dataset the model was asked for. + """ + from huggingface_hub.errors import HfHubHTTPError, LocalEntryNotFoundError + + try: + return download() + except (HfHubHTTPError, LocalEntryNotFoundError) as error: + raise FileNotFoundError( + f"Could not resolve the dataset {dataset_str!r}: {error}. " + "Check that the build id in the URI is published, or pass a " + "dataset= argument that is." + ) from error + + def _resolve_dataset_path(dataset_str): """Resolve a dataset string to a local file path, downloading if needed.""" if dataset_str.startswith("hf://datasets/"): @@ -248,11 +282,14 @@ def _resolve_dataset_path(dataset_str): version = None if "@" in repo_filename: repo_filename, version = repo_filename.rsplit("@", 1) - return hf_hub_download( - repo_id=f"{owner}/{repo}", - filename=repo_filename, - repo_type="dataset", - revision=version, + return _download_or_explain( + dataset_str, + lambda: hf_hub_download( + repo_id=f"{owner}/{repo}", + filename=repo_filename, + repo_type="dataset", + revision=version, + ), ) if "hf://" in dataset_str: from policyengine_core.tools.hugging_face import ( @@ -261,10 +298,13 @@ def _resolve_dataset_path(dataset_str): ) owner, repo, filename, version = parse_hf_url(dataset_str) - return download_huggingface_dataset( - repo=f"{owner}/{repo}", - repo_filename=filename, - version=version, + return _download_or_explain( + dataset_str, + lambda: download_huggingface_dataset( + repo=f"{owner}/{repo}", + repo_filename=filename, + version=version, + ), ) elif Path(dataset_str).exists(): return dataset_str @@ -289,7 +329,7 @@ def _is_hdfstore_format(file_path): return False -class Microsimulation(CoreMicrosimulation): +class Microsimulation(SPMSimulationMixin, CoreMicrosimulation): """ A microsimulation of the tax-benefit system for the United States, defined against the base microsimulation class in the -core package. @@ -317,6 +357,9 @@ def __init__(self, *args, **kwargs): start_instant: Annotated[str, "ISO date format YYYY-MM-DD"] = kwargs.pop( "start_instant", DEFAULT_START_DATE ) + args, kwargs = self._prepare_spm_system( + args, kwargs, kwargs.pop("spm", None), start_instant + ) dataset = kwargs.get("dataset") if dataset is None: diff --git a/policyengine_us/tests/code_health/test_batched_per_subdir_budget.py b/policyengine_us/tests/code_health/test_batched_per_subdir_budget.py new file mode 100644 index 00000000000..28436103caa --- /dev/null +++ b/policyengine_us/tests/code_health/test_batched_per_subdir_budget.py @@ -0,0 +1,76 @@ +"""Per-subdir batching must not pack more reform combos than one subprocess holds. + +The batcher's per-subdir mode gives each proposal folder its own subprocess. +A folder whose distinct reform combos exceed MAX_BATCH_COMBO_WEIGHT must be +packed by combo weight instead: policyengine-core caches one full system per +combo for the life of the subprocess, and congress/tlaib (weight 8.25) peaked +at 15.0 GB on the 16 GB CI runner at main, leaving no headroom; a branch run +of the same batch produced no output for 39 minutes before CI killed it. +""" + +import importlib.util +from pathlib import Path + +import yaml + +BATCHER = Path(__file__).resolve().parents[1] / "test_batched.py" +SPEC = importlib.util.spec_from_file_location("test_batched", BATCHER) +batched = importlib.util.module_from_spec(SPEC) +SPEC.loader.exec_module(batched) + + +def _write_cases(path: Path, reforms: list) -> None: + path.parent.mkdir(parents=True, exist_ok=True) + cases = [ + { + "name": f"{path.stem} {reform}", + "period": 2026, + "reforms": reform, + "input": {"people": {"person": {"age": 30}}}, + "output": {"age": 30}, + } + for reform in reforms + ] + path.write_text(yaml.safe_dump(cases)) + + +def test_light_subdir_stays_one_batch(tmp_path): + _write_cases(tmp_path / "light" / "a.yaml", ["reform_1"]) + _write_cases(tmp_path / "light" / "b.yaml", ["reform_1", "reform_2"]) + assert batched.subdir_batches(tmp_path / "light") == [[str(tmp_path / "light")]] + + +def test_heavy_subdir_is_packed_by_combo_weight(tmp_path): + heavy = tmp_path / "heavy" + _write_cases(heavy / "a.yaml", ["r1", "r2", "r3", "r4"]) + _write_cases(heavy / "b.yaml", ["r5", "r6"]) + _write_cases(heavy / "nested" / "c.yaml", ["r7"]) + _write_cases(heavy / "nested" / "d.yaml", ["r7", "r8"]) + batches = batched.subdir_batches(heavy) + assert batches == [ + [str(heavy / "a.yaml")], + [ + str(heavy / "b.yaml"), + str(heavy / "nested" / "c.yaml"), + str(heavy / "nested" / "d.yaml"), + ], + ] + for batch in batches: + combos: set = set() + for file in batch: + combos |= set(batched.file_reform_combos(Path(file))) + assert batched.combo_weight(frozenset(combos)) <= batched.MAX_BATCH_COMBO_WEIGHT + + +def test_per_subdir_mode_splits_only_heavy_folders(tmp_path): + _write_cases(tmp_path / "light" / "a.yaml", ["r1"]) + _write_cases(tmp_path / "heavy" / "a.yaml", ["r1", "r2", "r3", "r4"]) + _write_cases(tmp_path / "heavy" / "b.yaml", ["r5", "r6"]) + _write_cases(tmp_path / "root.yaml", ["r9"]) + batches = batched.split_into_batches(tmp_path, 2, mode="per-subdir") + assert batches == [ + [str(tmp_path / "heavy" / "a.yaml")], + [str(tmp_path / "heavy" / "b.yaml")], + [str(tmp_path / "light")], + [str(tmp_path / "root.yaml")], + ] diff --git a/policyengine_us/tests/core/test_behavioral_response_measurements.py b/policyengine_us/tests/core/test_behavioral_response_measurements.py index 5a22b7a2504..dd488b51a7d 100644 --- a/policyengine_us/tests/core/test_behavioral_response_measurements.py +++ b/policyengine_us/tests/core/test_behavioral_response_measurements.py @@ -530,7 +530,9 @@ def test_substitution_channel_is_live_end_to_end(): }, } - simulation = Simulation(situation=situation, reform=reform) + simulation = Simulation( + situation=situation, reform=reform, spm={"geography_kind": "national"} + ) response = simulation.calculate("labor_supply_behavioral_response", year) measurements = simulation._behavioral_response_measurements[year] baseline_mtr = np.asarray(measurements["baseline_mtr"]) @@ -581,7 +583,9 @@ def test_capital_gains_channel_is_live_end_to_end(): }, } - simulation = Simulation(situation=situation, reform=reform) + simulation = Simulation( + situation=situation, reform=reform, spm={"geography_kind": "national"} + ) response = simulation.calculate("capital_gains_behavioral_response", year) measurements = simulation._behavioral_response_measurements[year] baseline_mtr = np.asarray(measurements["baseline_capital_gains_mtr"]) diff --git a/policyengine_us/tests/core/test_local_employee_taxes.py b/policyengine_us/tests/core/test_local_employee_taxes.py index 1a3d20fb3b9..15bbfb1582d 100644 --- a/policyengine_us/tests/core/test_local_employee_taxes.py +++ b/policyengine_us/tests/core/test_local_employee_taxes.py @@ -12,6 +12,7 @@ def make_local_tax_simulation( state_code: str, period: str = PERIOD, *, + spm: dict | None = None, pa_philadelphia_wage_tax_taxable_wages: float = 0, pa_philadelphia_wage_tax_resident: bool = False, pa_philadelphia_wage_tax_reduced_rate_eligible: bool = False, @@ -25,6 +26,7 @@ def make_local_tax_simulation( ) -> Simulation: return Simulation( tax_benefit_system=SYSTEM, + spm=spm, situation={ "people": { "person": { @@ -277,9 +279,11 @@ def test_st_louis_credit_does_not_pool_across_people(): def test_local_taxes_feed_household_net_income(): - baseline = make_local_tax_simulation("PA") + spm = {"geography_kind": "national"} + baseline = make_local_tax_simulation("PA", spm=spm) reformed = make_local_tax_simulation( "PA", + spm=spm, pa_philadelphia_wage_tax_taxable_wages=WAGES, pa_philadelphia_wage_tax_resident=True, co_denver_employee_occupational_privilege_tax_months=12, diff --git a/policyengine_us/tests/core/test_pandas3_compatibility.py b/policyengine_us/tests/core/test_pandas3_compatibility.py index 6bf6c7d0d78..05438846dfe 100644 --- a/policyengine_us/tests/core/test_pandas3_compatibility.py +++ b/policyengine_us/tests/core/test_pandas3_compatibility.py @@ -113,8 +113,9 @@ def test_state_parameter_lookup(self): This exercises the VectorialParameterNodeAtInstant.__getitem__ fix that converts pandas StringArray to numpy array. """ - # Create a simulation with households in different states + # Keep state parameter lookups distinct while holding SPM geography fixed. sim = Simulation( + spm={"geography_kind": "national"}, situation={ "people": { "person1": {"age": {"2024": 30}}, @@ -130,7 +131,7 @@ def test_state_parameter_lookup(self): "state_code": {"2024": "NY"}, }, }, - } + }, ) # This calculation involves state-based parameter lookups diff --git a/policyengine_us/tests/core/test_payroll_contributions.py b/policyengine_us/tests/core/test_payroll_contributions.py index d2012604712..a4aee1cb138 100644 --- a/policyengine_us/tests/core/test_payroll_contributions.py +++ b/policyengine_us/tests/core/test_payroll_contributions.py @@ -20,6 +20,7 @@ def make_simulation( pre_tax_health_insurance_premiums: float = 0, tip_income: float = 0, traditional_401k_contributions: float = 0, + spm: dict | None = None, ) -> Simulation: household = { "members": ["person"], @@ -30,6 +31,7 @@ def make_simulation( return Simulation( tax_benefit_system=SYSTEM, + spm=spm, situation={ "people": { "person": { @@ -1201,8 +1203,9 @@ def test_employer_total_payroll_tax_aggregates_employer_inputs(): def test_employee_state_payroll_tax_flows_into_household_net_income(): - wa_sim = make_simulation("WA") - tx_sim = make_simulation("TX") + # Isolate the state payroll tax difference using the same SPM geography. + wa_sim = make_simulation("WA", spm={"geography_kind": "national"}) + tx_sim = make_simulation("TX", spm={"geography_kind": "national"}) wa_state_payroll_tax = calculate(wa_sim, "employee_state_payroll_tax") household_tax_difference = calculate( diff --git a/policyengine_us/tests/core/test_run_selective_tests.py b/policyengine_us/tests/core/test_run_selective_tests.py index c4248538104..cbdddc0757b 100644 --- a/policyengine_us/tests/core/test_run_selective_tests.py +++ b/policyengine_us/tests/core/test_run_selective_tests.py @@ -1,3 +1,6 @@ +from types import SimpleNamespace + +from policyengine_us.tests import run_selective_tests from policyengine_us.tests.run_selective_tests import SelectiveTestRunner @@ -168,3 +171,31 @@ def test_limit_test_paths_ignores_deleted_direct_tests(): assert deleted_test not in limited_paths assert existing_test in limited_paths + + +def _fake_pytest(returncodes): + def run(cmd, *args, **kwargs): + return SimpleNamespace(returncode=returncodes[cmd[-1]]) + + return run + + +def test_paths_that_collect_no_tests_do_not_fail_the_run(monkeypatch): + """A changed support module under tests/ collects nothing; that is not a failure.""" + runner = SelectiveTestRunner() + monkeypatch.setattr( + run_selective_tests.subprocess, + "run", + _fake_pytest({"tests/a_test.py": 0, "tests/populace_fixture.py": 5}), + ) + assert runner.run_tests({"tests/a_test.py", "tests/populace_fixture.py"}) == 0 + + +def test_real_failures_still_fail_beside_an_empty_path(monkeypatch): + runner = SelectiveTestRunner() + monkeypatch.setattr( + run_selective_tests.subprocess, + "run", + _fake_pytest({"tests/a_test.py": 1, "tests/populace_fixture.py": 5}), + ) + assert runner.run_tests({"tests/a_test.py", "tests/populace_fixture.py"}) == 1 diff --git a/policyengine_us/tests/core/test_spm_system.py b/policyengine_us/tests/core/test_spm_system.py new file mode 100644 index 00000000000..2b60bcd08ed --- /dev/null +++ b/policyengine_us/tests/core/test_spm_system.py @@ -0,0 +1,476 @@ +"""Construction and dataset contracts that require the Python simulation API.""" + +import hashlib +import json + +import numpy as np +import pandas as pd +import pytest +from policyengine_core.reforms import Reform +from policyengine_core.periods import YEAR +from policyengine_core.variables import Variable +from spm_calculator.errors import SPMInputError +from spm_calculator.policyengine_adapter import FORMULA_OWNED_INPUTS + +from policyengine_us import Microsimulation, Simulation +from policyengine_us.data.dataset_schema import USMultiYearDataset, USSingleYearDataset +from policyengine_us.entities import Person +from policyengine_us.spm import create_spm_provider +from policyengine_us.system import DEFAULT_DATASET, _resolve_dataset_path, system + + +def single_person_situation(): + return { + "people": {"person": {"age": {2024: 40}}}, + "households": { + "household": {"members": ["person"], "county_fips": {2024: "06037"}} + }, + } + + +class household_market_income(Variable): + # Partial reform variables inherit their entity, period and dtype from the + # baseline variable. Cloning must not call this class without that baseline. + def formula(household, period, parameters): + return household.filled_array(123) + + +class UserReform(Reform): + def apply(self): + self.update_variable(household_market_income) + self.neutralize_variable("income_tax") + + +@pytest.mark.parametrize("reform_wrapper", [False, True]) +def test_chosen_system_and_simulation_clones_preserve_user_reform(reform_wrapper): + chosen = system.clone() + if reform_wrapper: + chosen = UserReform(chosen) + else: + UserReform.apply(chosen) + original_config = chosen.spm_forecast_provider + situation = single_person_situation() + situation["people"]["person"]["employment_income"] = {2024: 50_000} + simulation = Simulation( + tax_benefit_system=chosen, + situation=situation, + spm={"geography_kind": "national"}, + ) + assert chosen.spm_forecast_provider is original_config + assert chosen.spm_forecast_provider.geography_kind == "county" + assert simulation.calculate("household_market_income", 2024)[0] == 123 + assert simulation.calculate("income_tax", 2024)[0] == 0 + assert simulation.tax_benefit_system.variables["income_tax"].is_neutralized + simulation.calculate("spm_unit_spm_threshold", 2024) + + clone = simulation.clone() + clone.delete_arrays("household_market_income") + assert clone.calculate("household_market_income", 2024)[0] == 123 + assert clone.tax_benefit_system.variables["income_tax"].is_neutralized + assert clone.calculate("income_tax", 2024)[0] == 0 + assert clone.spm_provenance() == simulation.spm_provenance() + clone.tax_benefit_system.variables["income_tax"].is_neutralized = False + clone.delete_arrays("income_tax") + assert clone.calculate("income_tax", 2024)[0] == 4_016 + assert simulation.calculate("income_tax", 2024)[0] == 0 + assert simulation.tax_benefit_system.variables["income_tax"].is_neutralized + assert chosen.variables["income_tax"].is_neutralized + + +@pytest.mark.parametrize("simulation_type", [Simulation, Microsimulation]) +def test_reform_baseline_and_clones_calculate_original_tax(simulation_type): + if simulation_type is Simulation: + situation = single_person_situation() + situation["people"]["person"]["employment_income"] = {2024: 50_000} + inputs = {"situation": situation} + else: + source = small_dataset() + source.person["age"] = [40, 40] + source.person["employment_income"] = [50_000, 50_000] + inputs = {"dataset": source} + config = {"geography_kind": "national"} + ordinary = simulation_type(**inputs, spm=config) + changed = simulation_type(**inputs, reform=UserReform, spm=config) + + np.testing.assert_array_equal(ordinary.calculate("income_tax", 2024), 4_016) + np.testing.assert_array_equal(changed.calculate("income_tax", 2024), 0) + np.testing.assert_array_equal(changed.baseline.calculate("income_tax", 2024), 4_016) + np.testing.assert_array_equal( + changed.baseline.calculate("household_market_income", 2024), 50_000 + ) + np.testing.assert_array_equal( + changed.calculate("household_market_income", 2024), 123 + ) + + for original, expected in ((changed, 0), (changed.baseline, 4_016)): + clone = original.clone() + np.testing.assert_array_equal(clone.calculate("income_tax", 2024), expected) + clone.delete_arrays("income_tax") + np.testing.assert_array_equal(clone.calculate("income_tax", 2024), expected) + for population in clone.populations.values(): + assert population.entity._tax_benefit_system is clone.tax_benefit_system + for name, holder in population._holders.items(): + assert holder.variable is clone.tax_benefit_system.variables[name] + assert holder.simulation is clone + assert clone.tax_benefit_system.simulation is clone + + +def _parameter_fingerprint(system): + """Digest every authored parameter value, so a shared-tree edit is visible.""" + digest = hashlib.sha256() + for parameter in system.parameters.get_descendants(): + values = getattr(parameter, "values_list", None) + if values is None: + continue + digest.update(parameter.name.encode()) + for value_at_instant in values: + digest.update( + f"|{value_at_instant.instant_str}={value_at_instant.value}".encode() + ) + return digest.hexdigest() + + +def test_ordinary_simulation_shares_default_policy_state(): + """A plain household simulation must not rebuild the shipped policy. + + Core's TaxBenefitSystem.clone() rebuilds the parameter tree node by node and + empties both at-instant caches, and this country deep-copies every variable + on top, so cloning per request would make each household API call rebuild + the at-instant tree for every period it touches. Nothing distinguishes an + ordinary simulation's policy from the shared instance's, so it shares the + tree and reuses the variable objects, and only its receipts and its own + variable registry are private. + """ + parameters_before = _parameter_fingerprint(system) + variables_before = {name: id(value) for name, value in system.variables.items()} + provider_before = system.spm_forecast_provider + + simulation = Simulation(situation=single_person_situation()) + policy = simulation.tax_benefit_system + + # Same parameter tree and the same warm at-instant caches: no clone ran. + assert policy is not system + assert policy.parameters is system.parameters + assert policy._parameters_at_instant_cache is system._parameters_at_instant_cache + # Entities stay bound to the shared instance, exactly as before this change. + assert policy.entities is system.entities + # A private registry, holding the shared instance's own variable objects. + assert policy.variables is not system.variables + assert ( + policy.variables["household_net_income"] + is system.variables["household_net_income"] + ) + rebound = { + name + for name, variable in policy.variables.items() + if variable is not system.variables[name] + } + # Only the structural reform re-applied at this simulation's start instant + # rebinds a name, and it rebinds in the private registry. + assert len(rebound) < 10, sorted(rebound) + + # Nothing reached the shared instance. + assert {name: id(value) for name, value in system.variables.items()} == ( + variables_before + ) + assert _parameter_fingerprint(system) == parameters_before + assert system.spm_forecast_provider is provider_before + + # Receipts remain private per simulation. + simulation.calculate("spm_unit_spm_threshold", 2024) + assert set(simulation.spm_provenance()["years"]) == {"2024"} + assert system.spm_forecast_provider.provenance()["years"] == {} + later = Simulation(situation=single_person_situation()) + assert later.spm_provenance()["years"] == {} + # The second simulation reads the at-instant tree the first one built. + assert system.parameters._at_instant_cache + assert ( + later.tax_benefit_system.parameters._at_instant_cache + is system.parameters._at_instant_cache + ) + + +def test_applying_reform_to_calculated_clone_preserves_original_tax(): + situation = single_person_situation() + situation["people"]["person"]["employment_income"] = {2024: 50_000} + original = Simulation(situation=situation) + assert original.calculate("income_tax", 2024)[0] == 4_016 + assert original.calculate("household_market_income", 2024)[0] == 50_000 + clone = original.clone() + branch = clone.get_branch("shared_policy") + isolated_branch = clone.get_branch("separate_policy", clone_system=True) + clone.apply_reform(UserReform) + assert clone.calculate("income_tax", 2024)[0] == 0 + assert clone.calculate("household_market_income", 2024)[0] == 123 + assert branch.calculate("income_tax", 2024)[0] == 0 + assert branch.calculate("household_market_income", 2024)[0] == 123 + assert ( + branch.get_holder("income_tax").variable + is branch.tax_benefit_system.variables["income_tax"] + ) + assert isolated_branch.calculate("income_tax", 2024)[0] == 4_016 + assert isolated_branch.calculate("household_market_income", 2024)[0] == 50_000 + clone.tax_benefit_system.variables["income_tax"].is_neutralized = False + clone.delete_arrays("income_tax") + assert clone.calculate("income_tax", 2024)[0] == 4_016 + assert original.calculate("income_tax", 2024)[0] == 4_016 + assert original.calculate("household_market_income", 2024)[0] == 50_000 + assert ( + clone.get_holder("income_tax").variable + is clone.tax_benefit_system.variables["income_tax"] + ) + + +class clone_only_income(Variable): + value_type = float + entity = Person + definition_period = YEAR + label = "Income input added only by the test reform" + + +class AddedInputReform(Reform): + def apply(self): + self.update_variable(clone_only_income) + + +def test_clone_can_add_a_variable_without_mutating_original(): + original = Simulation(situation=single_person_situation()) + clone = original.clone() + branch = clone.get_branch("shared_policy") + clone.apply_reform(AddedInputReform) + clone.set_input("clone_only_income", 2024, [123]) + assert clone.calculate("clone_only_income", 2024)[0] == 123 + branch.set_input("clone_only_income", 2024, [456]) + assert branch.calculate("clone_only_income", 2024)[0] == 456 + assert clone.calculate("clone_only_income", 2024)[0] == 123 + assert "clone_only_income" not in original.tax_benefit_system.variables + assert not any(key[0] == "clone_only_income" for key in original._user_input_keys) + + +def test_reform_only_input_does_not_leak_into_baseline_cache_invalidation(): + situation = single_person_situation() + situation["people"]["person"].update( + employment_income={2024: 50_000}, clone_only_income={2024: 123} + ) + simulation = Simulation(situation=situation, reform=AddedInputReform) + assert simulation.calculate("clone_only_income", 2024)[0] == 123 + baseline = simulation.baseline + assert "clone_only_income" not in baseline.input_variables + assert not any(key[0] == "clone_only_income" for key in baseline._user_input_keys) + baseline.apply_reform(()) + assert baseline.calculate("income_tax", 2024)[0] == 4_016 + assert simulation.calculate("clone_only_income", 2024)[0] == 123 + + +def test_new_simulation_from_calculated_system_starts_fresh_receipts(): + first = Simulation(situation=single_person_situation()) + first.calculate("spm_unit_spm_threshold", 2024) + second = Simulation(first.tax_benefit_system, situation=single_person_situation()) + assert second.spm_config == first.spm_config + assert second.spm_provenance()["years"] == {} + second.calculate("spm_unit_spm_threshold", 2025) + assert set(first.spm_provenance()["years"]) == {"2024"} + assert set(second.spm_provenance()["years"]) == {"2025"} + + +def test_positional_reform_baseline_uses_selected_spm_configuration(): + simulation = Simulation( + None, + None, + single_person_situation(), + None, + (), + spm={"geography_kind": "national"}, + ) + assert simulation.baseline.spm_config == simulation.spm_config + assert simulation.baseline.calculate("spm_unit_geographic_adjustment", 2024)[0] == 1 + assert simulation.spm_provenance()["years"] == {} + + +def test_shared_policy_branch_detaches_spm_receipts(): + simulation = Simulation(situation=single_person_situation()) + simulation.calculate("spm_unit_spm_threshold", 2024) + branch = simulation.get_branch("spm_receipt_test") + assert ( + branch.tax_benefit_system.variables is simulation.tax_benefit_system.variables + ) + assert ( + branch.tax_benefit_system.parameters is simulation.tax_benefit_system.parameters + ) + assert branch.tax_benefit_system.simulation is branch + assert branch.spm_provenance() == simulation.spm_provenance() + branch.calculate("spm_unit_spm_threshold", 2025) + assert set(simulation.spm_provenance()["years"]) == {"2024"} + assert set(branch.spm_provenance()["years"]) == {"2024", "2025"} + + +@pytest.mark.parametrize( + "config", + [ + {"forecast_path": "/tmp/forecast.json"}, + {"path": "/tmp/forecast.json"}, + {"as_of": "2022-01-01"}, + {"as_of": "2026-99-01"}, + {"geography_kind": "national", "geography_id": "06037"}, + {"geography_kind": "county", "geography_id": "06037"}, + {"geography_kind": "metro"}, + {"county_vintage": None}, + ], +) +def test_unsupported_public_configuration_fails(config): + with pytest.raises(ValueError): + create_spm_provider(config) + + +def test_as_of_config_and_receipts_are_serializable(): + simulation = Simulation( + situation=single_person_situation(), spm={"as_of": "2026-09-09"} + ) + simulation.calculate("spm_unit_spm_threshold", 2035) + assert json.loads(json.dumps(simulation.spm_config))["as_of"] == "2026-09-09" + assert set(json.loads(json.dumps(simulation.spm_provenance()))["years"]) == {"2035"} + + +class _HubResponse: + """The minimum a Hugging Face HTTP error reads off its response.""" + + headers = {} + request = None + + +def _hub_failures(): + from huggingface_hub.errors import ( + LocalEntryNotFoundError, + RevisionNotFoundError, + ) + + return [ + # What an unpublished build id actually produces: the Hub 404s and the + # download then reports that the file is not in the local cache either. + LocalEntryNotFoundError( + "An error happened while trying to locate the file on the Hub and " + "we cannot find the requested file" + ), + RevisionNotFoundError( + "404 Client Error: Revision Not Found", response=_HubResponse() + ), + ] + + +@pytest.mark.parametrize("failure", _hub_failures(), ids=["offline", "revision"]) +def test_unresolved_dataset_build_names_the_uri_it_could_not_resolve( + failure, monkeypatch +): + """A Hub failure must say which dataset the model was asked for. + + Both of these otherwise reach the caller with no mention of the URI, so an + unpublished build id reads as a broken installation. + """ + import huggingface_hub + + def missing(*args, **kwargs): + raise failure + + monkeypatch.setattr(huggingface_hub, "hf_hub_download", missing) + with pytest.raises(FileNotFoundError) as error: + _resolve_dataset_path(DEFAULT_DATASET) + message = str(error.value) + assert DEFAULT_DATASET in message + assert str(failure) in message + assert error.value.__cause__ is failure + + +@pytest.mark.parametrize("simulation_type", [Simulation, Microsimulation]) +def test_integer_county_column_requires_county_fips_instead_of_reporting_unavailable( + simulation_type, +): + """Legacy population files store the CPS within-state code as an integer.""" + source = small_dataset() + source.household["county_fips"] = [5, 1] + simulation = simulation_type(dataset=source) + with pytest.raises(SPMInputError) as error: + simulation.calculate("spm_unit_spm_threshold", 2024) + assert error.value.code == "SPM_GEOGRAPHY_REQUIRED" + assert "five-digit string" in str(error.value) + assert 'geography_kind="national"' in str(error.value) + # The same population computes once an SPM area is selected explicitly. + national = simulation_type(dataset=source, spm={"geography_kind": "national"}) + assert np.all(national.calculate("spm_unit_spm_threshold", 2024) > 0) + + +def small_dataset(): + groups = ("household", "tax_unit", "spm_unit", "family", "marital_unit") + return USSingleYearDataset( + person=pd.DataFrame( + { + "person_id": [1, 2], + "age": [16, 40], + "is_spm_independent_minor_role": [True, False], + **{f"person_{entity}_id": [1, 2] for entity in groups}, + } + ), + household=pd.DataFrame( + { + "household_id": [1, 2], + "county_fips": ["06037", "36061"], + "household_weight": [1.0, 1.0], + } + ), + **{ + entity: pd.DataFrame({f"{entity}_id": [1, 2]}) + for entity in groups + if entity != "household" + }, + time_period=2024, + ) + + +def test_positional_microsimulation_dataset_uses_country_interception(): + source = small_dataset() + simulation = Microsimulation(None, None, None, source) + assert isinstance(simulation.dataset, USMultiYearDataset) + assert simulation.calculate("spm_measurement_adults", 2024).tolist() == [1, 1] + + +@pytest.mark.parametrize("simulation_type", [Simulation, Microsimulation]) +def test_dataset_roles_receipts_and_formula_owned_rejection_preserve_source( + simulation_type, +): + source = small_dataset() + before = [frame.copy(deep=True) for frame in source.tables] + simulation = simulation_type(dataset=source) + simulation.set_input(variable_name="age", period=2024, value=[16, 40]) + assert simulation.calculate("spm_measurement_adults", 2024).tolist() == [1, 1] + assert simulation.calculate("spm_unit_count_adults", 2024).tolist() == [0, 1] + assert np.all(simulation.calculate("spm_unit_spm_threshold", 2024) > 0) + assert len(simulation.spm_provenance()["geographies"]) == 2 + json.dumps(simulation.spm_provenance()) + for name in FORMULA_OWNED_INPUTS: + with pytest.raises(ValueError, match="formula-owned SPM output"): + simulation.set_input(variable_name=name, period=2024, value=[1.0, 1.0]) + for original, current in zip(before, source.tables): + pd.testing.assert_frame_equal(original, current) + + +@pytest.mark.parametrize("dataset_format", ["single", "multi", "hdfstore"]) +def test_microsimulation_rejects_stored_measurements_without_mutating_dataset( + dataset_format, tmp_path +): + source = small_dataset() + source.spm_unit["spm_unit_spm_threshold"] = [100.0, 200.0] + before = [frame.copy(deep=True) for frame in source.tables] + if dataset_format == "multi": + dataset = USMultiYearDataset(datasets=[source]) + elif dataset_format == "hdfstore": + path = tmp_path / "spm_input.h5" + source.save(path) + content = path.read_bytes() + dataset = str(path) + else: + dataset = source + with pytest.raises(ValueError, match="formula-owned SPM output"): + Microsimulation(dataset=dataset) + for original, current in zip(before, source.tables): + pd.testing.assert_frame_equal(original, current) + if dataset_format == "hdfstore": + assert path.read_bytes() == content diff --git a/policyengine_us/tests/microsimulation/data/fixtures/test_extend_single_year_dataset.py b/policyengine_us/tests/microsimulation/data/fixtures/test_extend_single_year_dataset.py index ba99df62bc1..0439c6ab2be 100644 --- a/policyengine_us/tests/microsimulation/data/fixtures/test_extend_single_year_dataset.py +++ b/policyengine_us/tests/microsimulation/data/fixtures/test_extend_single_year_dataset.py @@ -262,7 +262,13 @@ def mock_super_init(self, *args, **kwargs): if captured is not None: captured[0] = ds self.dataset = ds - self.tax_benefit_system = system_module.system + self.tax_benefit_system = kwargs.get("tax_benefit_system", system_module.system) + self.populations = self.tax_benefit_system.instantiate_entities() + for population in self.populations.values(): + population.simulation = self + self._user_input_keys = set() + self.branches = {} + self.baseline = None self.is_over_dataset = True self.input_variables = [] self.get_holder = lambda name: MockHolder() diff --git a/policyengine_us/tests/microsimulation/populace_fixture.py b/policyengine_us/tests/microsimulation/populace_fixture.py new file mode 100644 index 00000000000..625d5680099 --- /dev/null +++ b/policyengine_us/tests/microsimulation/populace_fixture.py @@ -0,0 +1,39 @@ +"""A small, contract-compliant slice of the shipped population. + +Two microsimulations over the whole default dataset do not fit a CI runner, and +``subsample`` rebuilds the entire dataset as one wide frame before it shrinks +it. Slice whole households out of the shipped file instead: the result keeps the +observed county inputs and the source-backed SPM independence roles that the +input contract requires, which the legacy policyengine-us-data CPS files do not +carry. +""" + +from policyengine_us.data.dataset_schema import USSingleYearDataset +from policyengine_us.system import DEFAULT_DATASET, _resolve_dataset_path + + +GROUP_ENTITIES = ("tax_unit", "spm_unit", "family", "marital_unit") + + +def default_population_sample(households: int = 500) -> USSingleYearDataset: + """Return the first ``households`` whole households of the default build.""" + source = USSingleYearDataset(file_path=_resolve_dataset_path(DEFAULT_DATASET)) + kept = set(source.household["household_id"].iloc[:households]) + person = source.person[source.person["person_household_id"].isin(kept)] + tables = { + "person": person, + "household": source.household[source.household["household_id"].isin(kept)], + } + for entity in GROUP_ENTITIES: + frame = getattr(source, entity) + identifiers = set(person[f"person_{entity}_id"]) + tables[entity] = frame[frame[f"{entity}_id"].isin(identifiers)] + # A partially sliced group entity would misstate its own composition. + members = source.person[source.person[f"person_{entity}_id"].isin(identifiers)] + assert len(members) == len(person), ( + f"{entity} spans households outside the slice; widen the slice" + ) + return USSingleYearDataset( + time_period=int(source.time_period), + **{name: frame.reset_index(drop=True) for name, frame in tables.items()}, + ) diff --git a/policyengine_us/tests/microsimulation/test_future_dated_reform.py b/policyengine_us/tests/microsimulation/test_future_dated_reform.py index c088290ea14..f2438aea90d 100644 --- a/policyengine_us/tests/microsimulation/test_future_dated_reform.py +++ b/policyengine_us/tests/microsimulation/test_future_dated_reform.py @@ -8,7 +8,9 @@ import numpy as np -DATASET = "hf://policyengine/policyengine-us-data/cps_2023.h5" +from policyengine_us.tests.microsimulation.populace_fixture import ( + default_population_sample, +) def test_future_dated_reform_zero_impact_before_start(): @@ -28,12 +30,14 @@ def test_future_dated_reform_zero_impact_before_start(): country_id="us", ) - baseline = Microsimulation(dataset=DATASET) - reformed = Microsimulation(dataset=DATASET, reform=reform) - # Identical default seed (the dataset name) keeps the sampled households - # aligned between the two simulations. - baseline.subsample(1_000) - reformed.subsample(1_000) + # A slice of the shipped population, not the legacy policyengine-us-data + # CPS files: those supply neither source-backed SPM independence roles nor + # county FIPS codes, so resource outputs fail closed over them. Slicing + # whole households keeps both simulations on exactly the same population + # without subsampling two copies of the whole build. + dataset = default_population_sample() + baseline = Microsimulation(dataset=dataset) + reformed = Microsimulation(dataset=dataset, reform=reform) # The pre-start year must be the FIRST period calculated on each # simulation: computing an earlier period first carries those results @@ -48,3 +52,11 @@ def test_future_dated_reform_zero_impact_before_start(): np.asarray(baseline_net_income), err_msg=(f"Reform starting 2029 changed household net income in {year}"), ) + + # Guard against a vacuous pass: the same population must show the reform in + # the first year it starts, or equality above would prove nothing. + started = np.asarray(reformed.calc("household_net_income", period=2029)) + unreformed = np.asarray(baseline.calc("household_net_income", period=2029)) + assert (started != unreformed).any(), ( + "Reform starting 2029 changed no household net income in 2029" + ) diff --git a/policyengine_us/tests/microsimulation/test_microsim.py b/policyengine_us/tests/microsimulation/test_microsim.py index 331fc436bf5..7fa111a3754 100644 --- a/policyengine_us/tests/microsimulation/test_microsim.py +++ b/policyengine_us/tests/microsimulation/test_microsim.py @@ -1,46 +1,49 @@ import pytest -DATASETS = [ - "hf://policyengine/policyengine-us-data/cps_2023.h5", - "hf://policyengine/policyengine-us-data/enhanced_cps_2024.h5", -] +CPS_2023 = "hf://policyengine/policyengine-us-data/cps_2023.h5" +ENHANCED_CPS_2024 = "hf://policyengine/policyengine-us-data/enhanced_cps_2024.h5" -YEARS = list(range(2024, 2026)) +def test_legacy_cps_file_with_stored_measurements_is_rejected_at_load(): + """cps_2023 stores a formula-owned SPM output, so it cannot be loaded. -@pytest.fixture(scope="module", params=DATASETS) -def dataset_sim(request): - """One subsampled Microsimulation per dataset, shared across the - parametrized years (the tests only read via calc, so sharing is safe).""" + The loader rejects saved measurement outputs before any geography selection + could matter, which is the point of the contract: a population file supplies + primitive inputs, and observed Census outputs are retained under separate + report-only names. + """ from policyengine_us import Microsimulation - sim = Microsimulation(dataset=request.param) - sim.subsample(1_000) - return sim + with pytest.raises(ValueError, match="formula-owned SPM output"): + Microsimulation(dataset=CPS_2023) + +def test_legacy_enhanced_cps_lacks_source_backed_spm_independence_roles(): + """enhanced_cps_2024 loads, but cannot produce SPM measurements. -@pytest.mark.parametrize("year", YEARS) -def test_microsim_runs(dataset_sim, year: int): + Its county column does hold five-digit FIPS codes, so geography is not what + fails. 18 of its 43,134 SPM units are a lone 15-to-17-year-old carrying no + source-backed SPM independence role, so those units classify no measurement + adult, and every output that reaches the threshold - household net income, + benefits, poverty and marginal tax rates included - fails closed over the + file. Assert on the threshold itself: it raises straight off the + composition, without building the whole resource chain over 43,134 units. + """ import numpy as np + from spm_calculator.errors import SPMInputError - sim = dataset_sim - hnet = sim.calc("household_net_income", period=year) - assert not hnet.isna().any(), "Some households have NaN net income." - # Deciles are 1-10, with -1 for negative income. - DECILES = [ - "household_income_decile", - "spm_unit_income_decile", - "income_decile", - ] - for decile_var in DECILES: - decile = sim.calc(decile_var) - assert np.all(decile >= -1) and np.all(decile <= 10), ( - f"{decile_var} out of bounds." - ) + from policyengine_us import Microsimulation + + simulation = Microsimulation(dataset=ENHANCED_CPS_2024) + counties = np.asarray(simulation.calculate("county_fips", 2024)).astype(str) + assert np.all(np.char.str_len(counties) == 5) + assert np.all(np.char.isdigit(counties)) - # Check that the microsim calculates important variables as nonzero in current year. - for var in ["employment_income", "self_employment_income"]: - assert sim.calc(var, period=2024).sum() > 0, f"{var} is zero in 2024." + adults = np.asarray(simulation.calculate("spm_measurement_adults", 2024)) + assert (adults < 1).sum() == 18 + with pytest.raises(SPMInputError) as error: + simulation.calculate("spm_unit_spm_threshold", 2024) + assert error.value.code == "SPM_COMPOSITION_REQUIRED" def test_county_persists_across_periods(): @@ -50,8 +53,8 @@ def test_county_persists_across_periods(): should return the stored value regardless of the period requested, since county is a time-invariant geographic variable. - Uses the NYC dataset which has county explicitly stored (unlike cps_2023 - which doesn't store county). + Uses the NYC dataset, which stores the county enum itself rather than only + a county FIPS column. """ import numpy as np from policyengine_us import Microsimulation @@ -92,7 +95,13 @@ def test_county_persists_across_periods(): def test_default_dataset_loads_and_runs(): """The no-argument default (certified Populace build) resolves via the hf://datasets/ path and entity-level interception, and produces sane - aggregates.""" + aggregates. + + This is the only population file the model ships, so it carries the + society-wide coverage the legacy policyengine-us-data CPS files used to + provide here: net income across the extended years, in-range income + deciles, and nonzero earnings. + """ import numpy as np from policyengine_us import Microsimulation from policyengine_us.system import DEFAULT_DATASET @@ -103,9 +112,21 @@ def test_default_dataset_loads_and_runs(): sim = Microsimulation() # no dataset -> DEFAULT_DATASET (hf://datasets/...) sim.subsample(1_000) - for year in (2024, 2026): + for year in (2024, 2025, 2026): hnet = sim.calc("household_net_income", period=year) assert not hnet.isna().any(), f"NaN household net income in {year}." assert sim.calc("adjusted_gross_income", period=2026).sum() > 0, ( "Total AGI should be positive on the default dataset." ) + # Deciles are 1-10, with -1 for negative income. + for decile_variable in ( + "household_income_decile", + "spm_unit_income_decile", + "income_decile", + ): + decile = sim.calc(decile_variable) + assert np.all(decile >= -1) and np.all(decile <= 10), ( + f"{decile_variable} out of bounds." + ) + for variable in ("employment_income", "self_employment_income"): + assert sim.calc(variable, period=2024).sum() > 0, f"{variable} is zero in 2024." diff --git a/policyengine_us/tests/policy/baseline/household/income/spm_unit/spm_unit_capped_housing_subsidy.yaml b/policyengine_us/tests/policy/baseline/household/income/spm_unit/spm_unit_capped_housing_subsidy.yaml index 4052983dd9d..e1b3e5568b5 100644 --- a/policyengine_us/tests/policy/baseline/household/income/spm_unit/spm_unit_capped_housing_subsidy.yaml +++ b/policyengine_us/tests/policy/baseline/household/income/spm_unit/spm_unit_capped_housing_subsidy.yaml @@ -1,29 +1,81 @@ -- name: Housing subsidy is capped at the SPM housing need less tenant payment. +# Canonical spm-calculator 1.0.0 forecast spm-rolling-2026-09-09 gives the +# 2024 housing portion for a two-adult/two-child renter unit in county 06037: +# 39_219.893902 * 0.4305200747 * 1.6 = 27_015.922643864185. +# HUD inputs exercise the country benefit cap against this raw canonical amount. +- name: Case 1, housing subsidy is capped at housing need less tenant payment. period: 2024 + absolute_error_margin: 0.01 input: - spm_unit_spm_threshold_housing_portion: 17_720 - is_eligible_for_housing_assistance: true - hud_hap: 30_000 - hud_ttp: 10_000 + people: + person1: {age: 40} + person2: {age: 40} + person3: {age: 10} + person4: {age: 8} + spm_units: + spm_unit: + members: [person1, person2, person3, person4] + spm_unit_tenure_type: RENTER + is_eligible_for_housing_assistance: true + takes_up_housing_assistance_if_eligible: true + hud_hap: 30_000 + hud_ttp: 10_000 + households: + household: + members: [person1, person2, person3, person4] + state_code: CA + county_fips: "06037" output: - spm_unit_capped_housing_subsidy: 7_720 + housing_assistance: 30_000 + spm_unit_capped_housing_subsidy: 17_015.922643864185 -- name: Housing subsidy is not reduced when below the SPM cap. +- name: Case 2, housing subsidy below the SPM cap is retained. period: 2024 + absolute_error_margin: 0.01 input: - spm_unit_spm_threshold_housing_portion: 17_720 - is_eligible_for_housing_assistance: true - hud_hap: 5_000 - hud_ttp: 10_000 + people: + person1: {age: 40} + person2: {age: 40} + person3: {age: 10} + person4: {age: 8} + spm_units: + spm_unit: + members: [person1, person2, person3, person4] + spm_unit_tenure_type: RENTER + is_eligible_for_housing_assistance: true + takes_up_housing_assistance_if_eligible: true + hud_hap: 5_000 + hud_ttp: 10_000 + households: + household: + members: [person1, person2, person3, person4] + state_code: CA + county_fips: "06037" output: + housing_assistance: 5_000 spm_unit_capped_housing_subsidy: 5_000 -- name: Housing subsidy cap floors at zero. +- name: Case 3, tenant payment above canonical housing need floors the cap at zero. period: 2024 + absolute_error_margin: 0.01 input: - spm_unit_spm_threshold_housing_portion: 4_000 - is_eligible_for_housing_assistance: true - hud_hap: 5_000 - hud_ttp: 5_000 + people: + person1: {age: 40} + person2: {age: 40} + person3: {age: 10} + person4: {age: 8} + spm_units: + spm_unit: + members: [person1, person2, person3, person4] + spm_unit_tenure_type: RENTER + is_eligible_for_housing_assistance: true + takes_up_housing_assistance_if_eligible: true + hud_hap: 5_000 + hud_ttp: 30_000 + households: + household: + members: [person1, person2, person3, person4] + state_code: CA + county_fips: "06037" output: + housing_assistance: 5_000 spm_unit_capped_housing_subsidy: 0 diff --git a/policyengine_us/tests/policy/baseline/household/income/spm_unit/spm_unit_net_income.yaml b/policyengine_us/tests/policy/baseline/household/income/spm_unit/spm_unit_net_income.yaml index 083d5aa0f32..1fb105a5c2d 100644 --- a/policyengine_us/tests/policy/baseline/household/income/spm_unit/spm_unit_net_income.yaml +++ b/policyengine_us/tests/policy/baseline/household/income/spm_unit/spm_unit_net_income.yaml @@ -23,6 +23,7 @@ period: 2024 input: age: 30 + county_fips: "06037" spm_unit_market_income: 0 child_support_received: 1 workers_compensation: 2 diff --git a/policyengine_us/tests/policy/baseline/household/income/spm_unit/spm_unit_spm_threshold_housing_portion.yaml b/policyengine_us/tests/policy/baseline/household/income/spm_unit/spm_unit_spm_threshold_housing_portion.yaml index e998db6d4a0..e52944277d7 100644 --- a/policyengine_us/tests/policy/baseline/household/income/spm_unit/spm_unit_spm_threshold_housing_portion.yaml +++ b/policyengine_us/tests/policy/baseline/household/income/spm_unit/spm_unit_spm_threshold_housing_portion.yaml @@ -1,13 +1,70 @@ -- name: Housing portion isolates the geographically adjusted housing share. +# Canonical spm-calculator 1.0.0 forecast spm-rolling-2026-09-09, 2024: +# County 06037 resolves to Los Angeles metro 31080, with rent index 1.6. +# Two adults and two children have equivalence factor 1. The canonical housing +# amounts below use the published reference threshold * housing share * 1.6, +# before PolicyEngine casts the final amount to its variable storage dtype. +- name: Case 1, renter housing portion uses the canonical county amount. period: 2024 absolute_error_margin: 0.01 input: - # Renter housing share is 0.443. A 1.5 local-rent ratio implies: - # geoadj = 1 - 0.443 + 0.443 * 1.5 = 1.2215. - # For a $40,000 unadjusted threshold, the adjusted threshold is $48,860 - # and the adjusted housing portion is $40,000 * 0.443 * 1.5 = $26,580. - spm_unit_unadjusted_spm_threshold: 40_000 - spm_unit_geographic_adjustment: 1.2215 - spm_unit_tenure_type: RENTER + people: + person1: {age: 40} + person2: {age: 40} + person3: {age: 10} + person4: {age: 8} + spm_units: + spm_unit: + members: [person1, person2, person3, person4] + spm_unit_tenure_type: RENTER + households: + household: + members: [person1, person2, person3, person4] + state_code: CA + county_fips: "06037" output: - spm_unit_spm_threshold_housing_portion: 26_580 + # 39_219.893902 * 0.4305200747 * 1.6. + spm_unit_spm_threshold_housing_portion: 27_015.922643864185 + +- name: Case 2, mortgaged owner housing portion uses the canonical county amount. + period: 2024 + absolute_error_margin: 0.01 + input: + people: + person1: {age: 40} + person2: {age: 40} + person3: {age: 10} + person4: {age: 8} + spm_units: + spm_unit: + members: [person1, person2, person3, person4] + spm_unit_tenure_type: OWNER_WITH_MORTGAGE + households: + household: + members: [person1, person2, person3, person4] + state_code: CA + county_fips: "06037" + output: + # 39_230.994457 * 0.4306812111 * 1.6. + spm_unit_spm_threshold_housing_portion: 27_033.683528637033 + +- name: Case 3, outright owner housing portion uses the canonical county amount. + period: 2024 + absolute_error_margin: 0.01 + input: + people: + person1: {age: 40} + person2: {age: 40} + person3: {age: 10} + person4: {age: 8} + spm_units: + spm_unit: + members: [person1, person2, person3, person4] + spm_unit_tenure_type: OWNER_WITHOUT_MORTGAGE + households: + household: + members: [person1, person2, person3, person4] + state_code: CA + county_fips: "06037" + output: + # 32_878.594848 * 0.3206844041 * 1.6. + spm_unit_spm_threshold_housing_portion: 16_869.84415436194 diff --git a/policyengine_us/tests/policy/baseline/household/income/spm_unit/test_spm_unit_spm_threshold.py b/policyengine_us/tests/policy/baseline/household/income/spm_unit/test_spm_unit_spm_threshold.py index 012c7b647f8..4849d9f7086 100644 --- a/policyengine_us/tests/policy/baseline/household/income/spm_unit/test_spm_unit_spm_threshold.py +++ b/policyengine_us/tests/policy/baseline/household/income/spm_unit/test_spm_unit_spm_threshold.py @@ -1,261 +1,255 @@ -"""Regression tests for the spm_unit_spm_threshold formula. - -Verifies that: -1. Historical published years (2015-2024) match the Census Bureau's - published Betson reference thresholds exactly (via spm-calculator). -2. Post-published years uprate via PolicyEngine's ``gov.bls.cpi.cpi_u`` - parameter. -3. Composition and tenure changes between periods flow through to the - threshold while applying the unit-specific geographic adjustment. -4. The Betson three-parameter equivalence scale is applied (a 2A2C - reference family at the renter national base equals 39430 in 2024). +"""Canonical forecast selection and country-construction contract regressions. + +These Python cases exercise configuration, errors, registry timing and clone +receipts, which the YAML input schema cannot express. """ +import json + import numpy as np +import pandas as pd import pytest + from policyengine_us import Simulation -from policyengine_us.variables.household.income.spm_unit.spm_unit_reference_spm_threshold import ( - LATEST_PUBLISHED_SPM_THRESHOLD_YEAR, - _reference_threshold_array, -) -from policyengine_us.variables.household.income.spm_unit.spm_unit_tenure_type import ( - SPMUnitTenureType, -) -from spm_calculator.equivalence_scale import spm_equivalence_scale -from spm_calculator.forecast import HISTORICAL_THRESHOLDS -from spm_calculator.geoadj import get_cd_geoadj, get_housing_share - - -def _cpi_u(): - from policyengine_us import CountryTaxBenefitSystem - - return CountryTaxBenefitSystem().parameters.gov.bls.cpi.cpi_u - - -def test_reference_threshold_matches_census_for_published_years(): - """The per-tenure reference threshold must equal the Census BLS - published value for every year spm-calculator has on record.""" - cpi_u = _cpi_u() - for year, expected in HISTORICAL_THRESHOLDS.items(): - for tenure_enum, key in [ - (SPMUnitTenureType.RENTER, "renter"), - (SPMUnitTenureType.OWNER_WITH_MORTGAGE, "owner_with_mortgage"), - ( - SPMUnitTenureType.OWNER_WITHOUT_MORTGAGE, - "owner_without_mortgage", - ), - ]: - got = _reference_threshold_array( - np.array([tenure_enum]), - year, - cpi_u, - )[0] - assert got == expected[key], ( - f"{year} {key}: got {got}, expected {expected[key]}" - ) - - -def test_reference_threshold_uprates_with_cpi_u_past_latest_published(): - """Post-2024 reference threshold must equal the latest published - value scaled by PolicyEngine's CPI-U ratio.""" - cpi_u = _cpi_u() - latest = LATEST_PUBLISHED_SPM_THRESHOLD_YEAR - target_year = latest + 2 - factor = float(cpi_u(f"{target_year}-02-01") / cpi_u(f"{latest}-02-01")) - published_renter = HISTORICAL_THRESHOLDS[latest]["renter"] - got = _reference_threshold_array( - np.array([SPMUnitTenureType.RENTER]), - target_year, - cpi_u, - )[0] - assert got == published_renter * factor - - -def test_formula_respects_composition_change_between_years(): - """When a child ages into adulthood between the prior and current - period, the threshold must rescale by the Betson equivalence-scale - ratio.""" - YEAR = 2026 - PRIOR = 2025 - - # Build a simulation where person "teen" is 17 (child) in PRIOR and - # 18 (adult) in YEAR, converting the SPM unit from 2A2C to 3A1C. - cpi_u = _cpi_u() - current_base = _reference_threshold_array( - np.array([SPMUnitTenureType.RENTER]), - YEAR, - cpi_u, - )[0] - current_equiv = spm_equivalence_scale(3, 1) # teen is adult in YEAR - - situation = { - "people": { - "adult1": {"age": {PRIOR: 40, YEAR: 41}}, - "adult2": {"age": {PRIOR: 40, YEAR: 41}}, - "teen": {"age": {PRIOR: 17, YEAR: 18}}, - "child": {"age": {PRIOR: 5, YEAR: 6}}, - }, +from policyengine_us.data.dataset_schema import USSingleYearDataset +from policyengine_us.spm import create_spm_provider +from spm_calculator.errors import SPMInputError +from spm_calculator.release import SPMUnit +from spm_calculator.rolling_forecast import load_forecast + + +AMOUNTS = { + "spm_unit_reference_spm_threshold": "reference_threshold", + "spm_unit_unadjusted_spm_threshold": "unadjusted_threshold", + "spm_unit_geographic_adjustment": "geographic_factor", + "spm_unit_spm_threshold": "threshold", + "spm_unit_spm_threshold_housing_portion": "housing_portion", +} + + +def household(*, county=None, tenure="RENTER", people=None): + people = people or { + "person1": {"age": {2024: 40}}, + "person2": {"age": {2024: 40}}, + "person3": {"age": {2024: 10}}, + "person4": {"age": {2024: 8}}, + } + members = list(people) + location = {"members": members, "state_code": {2024: "CA"}} + if county is not None: + location["county_fips"] = {2024: county} + return { + "people": people, + "households": {"household": location}, "spm_units": { "spm_unit": { - "members": ["adult1", "adult2", "teen", "child"], - "spm_unit_tenure_type": { - PRIOR: "RENTER", - YEAR: "RENTER", - }, + "members": members, + "spm_unit_tenure_type": {2024: tenure}, } }, } - sim = Simulation(situation=situation) - expected = current_base * current_equiv - got_reference = float(sim.calculate("spm_unit_reference_spm_threshold", YEAR)[0]) - got_unadjusted = float(sim.calculate("spm_unit_unadjusted_spm_threshold", YEAR)[0]) - got = float(sim.calculate("spm_unit_spm_threshold", YEAR)[0]) - assert got_reference == pytest.approx(current_base, rel=1e-4) - assert got_unadjusted == pytest.approx(expected, rel=1e-4) - assert got == pytest.approx(expected, rel=1e-4) - - -def test_formula_uses_congressional_district_geographic_adjustment(): - """The geographic adjustment is calculated from raw household geography - and SPM tenure.""" - YEAR = 2023 - CD_GEOID = 101 - - cpi_u = _cpi_u() - base = _reference_threshold_array( - np.array([SPMUnitTenureType.OWNER_WITH_MORTGAGE]), - YEAR, - cpi_u, - )[0] - equiv = spm_equivalence_scale(2, 2) - geoadj = get_cd_geoadj( - CD_GEOID, - year=2023, - tenure="owner_with_mortgage", + +def canonical(forecast, year, tenure, *, adults=2, children=2, county=None): + location = ( + forecast.resolve_county(year, county) + if county is not None + else {"kind": "national", "area_id": None} + ) + return forecast.calculate_unit( + SPMUnit( + unit_id="test", + year=year, + num_adults=adults, + num_children=children, + tenure=tenure.lower(), + geography_kind=location["kind"], + geography_id=location["area_id"], + ) ) - situation = { - "people": { - "a1": {"age": {YEAR: 40}}, - "a2": {"age": {YEAR: 40}}, - "k1": {"age": {YEAR: 5}}, - "k2": {"age": {YEAR: 3}}, - }, - "households": { - "household": { - "members": ["a1", "a2", "k1", "k2"], - "congressional_district_geoid": { - YEAR: CD_GEOID, - }, - } - }, - "spm_units": { - "spm_unit": { - "members": ["a1", "a2", "k1", "k2"], - "spm_unit_tenure_type": { - YEAR: "OWNER_WITH_MORTGAGE", - }, - } - }, + +@pytest.mark.parametrize( + "tenure", ["RENTER", "OWNER_WITH_MORTGAGE", "OWNER_WITHOUT_MORTGAGE"] +) +def test_every_amount_uses_canonical_final_value_for_every_supported_year(tenure): + forecast = load_forecast() + assert list(forecast.years) == list(range(2022, 2036)) + # Explicit period inputs cover the artifact, including its historical years. + situation = household(county="06037", tenure=tenure) + for person in situation["people"].values(): + person["age"] = { + year: next(iter(person["age"].values())) for year in forecast.years + } + situation["households"]["household"]["county_fips"] = { + year: "06037" for year in forecast.years } + situation["spm_units"]["spm_unit"]["spm_unit_tenure_type"] = { + year: tenure for year in forecast.years + } + simulation = Simulation(situation=situation) + for year in forecast.years: + expected = canonical(forecast, year, tenure, county="06037") + for variable, field in AMOUNTS.items(): + result = simulation.calculate(variable, year) + assert result[0] == np.asarray(expected[field], dtype=result.dtype) + assert set(simulation.spm_provenance()["years"]) == { + str(year) for year in forecast.years + } + - sim = Simulation(situation=situation) - got = float(sim.calculate("spm_unit_spm_threshold", YEAR)[0]) - got_unadjusted = float(sim.calculate("spm_unit_unadjusted_spm_threshold", YEAR)[0]) - got_housing_portion = float( - sim.calculate("spm_unit_spm_threshold_housing_portion", YEAR)[0] +def test_published_national_reference_is_exact_before_storage_cast(): + # 2024 Census/BLS Betson renter amount in the pinned official workbook. + provider = create_spm_provider({"geography_kind": "national"}) + result = provider.calculate_unit(year=2024, adults=2, children=2, tenure="renter") + assert result["reference_threshold"] == 39_219.893902 + + +@pytest.mark.parametrize( + "county,code", + [ + (None, "SPM_GEOGRAPHY_REQUIRED"), + ("", "SPM_GEOGRAPHY_REQUIRED"), + # County FIPS is a string input, so every non-string value reaches the + # forecast as a string. A CPS within-state code, a FIPS code that lost + # its leading zero to an integer column and a missing value are absent + # counties, not unavailable ones. + (5, "SPM_GEOGRAPHY_REQUIRED"), + (6037, "SPM_GEOGRAPHY_REQUIRED"), + ("6037", "SPM_GEOGRAPHY_REQUIRED"), + (float("nan"), "SPM_GEOGRAPHY_REQUIRED"), + ("99999", "SPM_GEOGRAPHY_UNAVAILABLE"), + ], +) +def test_state_only_or_unknown_county_never_selects_national(county, code): + simulation = Simulation(situation=household(county=county)) + with pytest.raises(SPMInputError) as error: + simulation.calculate("spm_unit_spm_threshold", 2025) + assert error.value.to_dict()["code"] == code + assert simulation.spm_config["geography_kind"] == "county" + if code == "SPM_GEOGRAPHY_REQUIRED": + # The message has to name the caller's fix, not the artifact. + message = str(error.value) + assert "five-digit string" in message + assert 'geography_kind="national"' in message + + +def test_explicit_national_and_serialized_round_trip(): + simulation = Simulation(situation=household(), spm={"geography_kind": "national"}) + config = json.loads(json.dumps(simulation.spm_config)) + assert config["forecast_content_sha256"] == load_forecast().content_sha256 + assert config["scenario"] == load_forecast().default_scenario + repeated = Simulation(situation=household(), spm=config) + assert ( + repeated.calculate("spm_unit_spm_threshold", 2025)[0] + == simulation.calculate("spm_unit_spm_threshold", 2025)[0] ) - unadjusted = base * equiv - housing_share = get_housing_share("owner_with_mortgage") - expected_housing_portion = unadjusted * (geoadj + housing_share - 1) - expected = unadjusted * geoadj - assert got_unadjusted == pytest.approx(unadjusted, rel=1e-5) - assert got_housing_portion == pytest.approx( - expected_housing_portion, - rel=1e-5, + assert simulation.calculate("spm_unit_geographic_adjustment", 2025)[0] == 1 + config["scenario"] = "unavailable" + assert simulation.spm_config["scenario"] != "unavailable" + + +@pytest.mark.parametrize( + "config", + [ + {"forecast_content_sha256": "0" * 64}, + {"scenario": "unavailable"}, + {"geography_kind": "state"}, + {"missing_geography": "national"}, + {"year_policy": "pe_cpi_u"}, + ], +) +def test_unavailable_selection_or_consumer_extrapolation_rejected(config): + with pytest.raises(ValueError): + create_spm_provider(config) + + +def test_unknown_year_has_no_country_cpi_extrapolation(): + simulation = Simulation(situation=household(county="06037")) + with pytest.raises(ValueError, match="no entry"): + simulation.calculate("spm_unit_spm_threshold", 2099) + + +@pytest.mark.parametrize( + "role", + ["is_household_head", "is_household_spouse", "is_spm_independent_minor_role"], +) +def test_minor_primitive_exists_before_input_parsing_and_does_not_change_benefit_counts( + role, +): + simulation = Simulation( + situation=household(people={"person1": {"age": {2024: 16}, role: True}}), + spm={"geography_kind": "national"}, ) - assert got == pytest.approx(expected, rel=1e-5) - - -def test_geographic_adjustment_defaults_to_one_without_congressional_district(): - """With no congressional district input, the threshold uses the national - reference threshold.""" - YEAR = 2024 - - cpi_u = _cpi_u() - base = _reference_threshold_array( - np.array([SPMUnitTenureType.RENTER]), - YEAR, - cpi_u, - )[0] - equiv = spm_equivalence_scale(2, 2) - - situation = { - "people": { - "a1": {"age": {YEAR: 40}}, - "a2": {"age": {YEAR: 40}}, - "k1": {"age": {YEAR: 5}}, - "k2": {"age": {YEAR: 3}}, - }, - "spm_units": { - "spm_unit": { - "members": ["a1", "a2", "k1", "k2"], - "spm_unit_tenure_type": { - YEAR: "RENTER", - }, - } - }, - } + assert simulation.calculate("spm_measurement_adults", 2025)[0] == 1 + assert simulation.calculate("spm_measurement_children", 2025)[0] == 0 + assert simulation.calculate("spm_unit_count_adults", 2025)[0] == 0 + assert simulation.calculate("spm_unit_count_children", 2025)[0] == 1 + assert simulation.calculate("spm_unit_spm_threshold", 2025)[0] > 0 - sim = Simulation(situation=situation) - got = float(sim.calculate("spm_unit_spm_threshold", YEAR)[0]) - expected = base * equiv - assert got == pytest.approx(expected, rel=1e-5) - - -def test_prior_threshold_does_not_imply_geographic_adjustment(): - """Stored prior-year thresholds should not back out geographic adjustment; - geographic adjustment is formulaic from current raw geography.""" - YEAR = 2026 - PRIOR = 2025 - PRIOR_GEOADJ = 1.5 - - cpi_u = _cpi_u() - prior_base = _reference_threshold_array( - np.array([SPMUnitTenureType.RENTER]), - PRIOR, - cpi_u, - )[0] - current_base = _reference_threshold_array( - np.array([SPMUnitTenureType.RENTER]), - YEAR, - cpi_u, - )[0] - equiv = spm_equivalence_scale(2, 2) - - situation = { - "people": { - "a1": {"age": {PRIOR: 40, YEAR: 41}}, - "a2": {"age": {PRIOR: 40, YEAR: 41}}, - "k1": {"age": {PRIOR: 5, YEAR: 6}}, - "k2": {"age": {PRIOR: 3, YEAR: 4}}, - }, - "spm_units": { - "spm_unit": { - "members": ["a1", "a2", "k1", "k2"], - "spm_unit_spm_threshold": { - PRIOR: float(prior_base * equiv * PRIOR_GEOADJ), - }, - "spm_unit_tenure_type": { - PRIOR: "RENTER", - YEAR: "RENTER", - }, - } - }, - } - sim = Simulation(situation=situation) - got = float(sim.calculate("spm_unit_spm_threshold", YEAR)[0]) - expected = current_base * equiv - assert got == pytest.approx(expected, rel=1e-5) +def test_no_adult_requires_explicit_composition_instead_of_minimum_one(): + simulation = Simulation( + situation=household(people={"person1": {"age": {2024: 16}}}), + spm={"geography_kind": "national"}, + ) + with pytest.raises(SPMInputError) as error: + simulation.calculate("spm_unit_spm_threshold", 2025) + assert error.value.to_dict()["code"] == "SPM_COMPOSITION_REQUIRED" + + +def test_clone_keeps_cached_calculation_provenance_but_detaches_receipts(): + simulation = Simulation(situation=household(county="06037")) + expected = simulation.calculate("spm_unit_spm_threshold", 2025)[0] + clone = simulation.clone() + assert clone.calculate("spm_unit_spm_threshold", 2025)[0] == expected + assert clone.spm_provenance()["years"] == simulation.spm_provenance()["years"] + clone.calculate("spm_unit_spm_threshold", 2026) + assert "2026" not in simulation.spm_provenance()["years"] + detached = clone.spm_provenance() + detached["years"].clear() + assert clone.spm_provenance()["years"] + + +def test_policy_reform_baseline_uses_same_explicit_spm_selection(): + simulation = Simulation( + situation=household(), reform=(), spm={"geography_kind": "national"} + ) + assert simulation.baseline.spm_config == simulation.spm_config + assert ( + simulation.baseline.calculate("spm_unit_spm_threshold", 2025)[0] + == simulation.calculate("spm_unit_spm_threshold", 2025)[0] + ) + + +def test_dataset_cannot_override_formula_owned_threshold(): + person = pd.DataFrame( + { + "person_id": [1], + "age": [40], + **{ + f"person_{entity}_id": [1] + for entity in ( + "household", + "tax_unit", + "spm_unit", + "family", + "marital_unit", + ) + }, + } + ) + dataset = USSingleYearDataset( + person=person, + household=pd.DataFrame({"household_id": [1], "county_fips": ["06037"]}), + tax_unit=pd.DataFrame({"tax_unit_id": [1]}), + spm_unit=pd.DataFrame({"spm_unit_id": [1], "spm_unit_spm_threshold": [1.0]}), + family=pd.DataFrame({"family_id": [1]}), + marital_unit=pd.DataFrame({"marital_unit_id": [1]}), + time_period=2024, + ) + with pytest.raises( + ValueError, + match="Dataset supplies formula-owned SPM output spm_unit_spm_threshold", + ): + Simulation(dataset=dataset) diff --git a/policyengine_us/tests/run_selective_tests.py b/policyengine_us/tests/run_selective_tests.py index e3fbf8c784f..b41738d6295 100644 --- a/policyengine_us/tests/run_selective_tests.py +++ b/policyengine_us/tests/run_selective_tests.py @@ -15,6 +15,8 @@ SOURCE_FILE_SUFFIXES = (".py", ".yaml", ".yml") TEST_INFRASTRUCTURE_FILES = ("run_selective_tests.py", "test_batched.py") +# pytest exits 5 when a path collects no tests (see pytest.ExitCode). +PYTEST_NO_TESTS_COLLECTED = 5 CRITICAL_TEST_TRIGGER_FILES = ("pyproject.toml", "requirements", "Makefile") STOP_TEST_DIRS = frozenset( { @@ -519,6 +521,11 @@ def run_tests( cmd = base_cmd + [test_path] print(f"\nRunning command: {' '.join(cmd)}") result = subprocess.run(cmd) + if result.returncode == PYTEST_NO_TESTS_COLLECTED: + # A changed file under tests/ with nothing to collect is a + # support module (a fixture or helper), not a failing test. + print(f"No tests collected in {test_path}; nothing to run there.") + continue if result.returncode != 0: worst_returncode = result.returncode diff --git a/policyengine_us/tests/test_batched.py b/policyengine_us/tests/test_batched.py index 9bb39ae551d..bf84fbee3f3 100644 --- a/policyengine_us/tests/test_batched.py +++ b/policyengine_us/tests/test_batched.py @@ -166,6 +166,27 @@ def _format_rss(peak_rss_mb: Optional[float]) -> str: return f"{peak_rss_mb:.0f} MB" if peak_rss_mb is not None else "n/a" +def subdir_batches(subdir: Path) -> List[List[str]]: + """One batch per subdir unless its distinct combos exceed the budget. + + A proposal folder normally shares one subprocess so files that request + the same reform share one cached system. Once the folder's DISTINCT + combos exceed MAX_BATCH_COMBO_WEIGHT that shared subprocess is exactly + what exhausts the runner: congress/tlaib carries weight 8.25 and peaked + at 15.0 GB on the 16 GB runner at main (CI run 34637166889), leaving no + headroom, and a branch run of the same batch produced no output for 39 + minutes before CI killed it. Pack such a folder's files by combo weight + instead. + """ + files = sorted(subdir.rglob("*.yaml")) + combos: set = set() + for file in files: + combos |= set(file_reform_combos(file)) + if combo_weight(frozenset(combos)) <= MAX_BATCH_COMBO_WEIGHT: + return [[str(subdir)]] + return pack_files_by_combo_weight(files) + + def split_into_batches( base_path: Path, num_batches: int, @@ -183,7 +204,8 @@ def split_into_batches( exclude: List of directory names to exclude (for contrib tests) mode: Batching mode. "auto" (default) uses the per-path heuristics below. "per-subdir" runs each immediate subdir as its own batch - with loose yamls collected into a trailing batch. "per-file" + (packed by reform-combo weight when it exceeds the budget) with + loose yamls collected into a trailing batch. "per-file" runs every yaml (recursively) as its own batch. """ if exclude is None: @@ -201,7 +223,7 @@ def split_into_batches( if item.is_dir() and item.name not in exclude ) root_files = sorted(base_path.glob("*.yaml")) - batches = [[str(s)] for s in subdirs] + batches = [batch for s in subdirs for batch in subdir_batches(s)] if root_files: batches.append([str(f) for f in root_files]) return batches @@ -783,7 +805,7 @@ def main(): "--mode", choices=["auto", "per-subdir", "per-file"], default="auto", - help="Batching mode. 'per-subdir' = each immediate subdir is its own batch; 'per-file' = each yaml is its own batch.", + help="Batching mode. 'per-subdir' = each immediate subdir is its own batch, split by reform-combo weight when over budget; 'per-file' = each yaml is its own batch.", ) parser.add_argument( "--workers", diff --git a/policyengine_us/tests/unit/test_spm_integration_contract.py b/policyengine_us/tests/unit/test_spm_integration_contract.py new file mode 100644 index 00000000000..e03e00f18a1 --- /dev/null +++ b/policyengine_us/tests/unit/test_spm_integration_contract.py @@ -0,0 +1,421 @@ +"""Cross-package contracts requiring configuration and dtype-aware assertions. + +These cases exercise the public Simulation constructor and computed HUD benefits; +the YAML runner cannot express forecast configuration or compare storage casts. +""" + +import json + +import numpy as np +import pytest + +from policyengine_us import Simulation +from spm_calculator.errors import SPMInputError +from spm_calculator.release import SPMUnit +from spm_calculator.rolling_forecast import load_forecast + + +YEAR = 2024 + + +def household(*, people=None, county=None, earnings=0, rent=36_000): + if people is None: + people = { + "person1": { + "age": {YEAR: 40}, + "employment_income": {YEAR: earnings}, + "pre_subsidy_rent": {YEAR: rent}, + } + } + members = list(people) + location = { + "members": members, + "state_code": {YEAR: "CA"}, + "pha_payment_standard": {YEAR: 36_000}, + } + if county is not None: + location["county_fips"] = {YEAR: county} + return { + "people": people, + "households": {"household": location}, + "spm_units": { + "spm_unit": { + "members": members, + "spm_unit_tenure_type": {YEAR: "RENTER"}, + "receives_housing_assistance": {YEAR: True}, + } + }, + } + + +@pytest.fixture(scope="module") +def forecast(): + return load_forecast() + + +def canonical(forecast, *, adults=1, children=0, geography_kind="national", area=None): + return forecast.calculate_unit( + SPMUnit( + unit_id="test", + year=YEAR, + num_adults=adults, + num_children=children, + tenure="renter", + geography_kind=geography_kind, + geography_id=area, + ) + ) + + +@pytest.mark.parametrize( + "earnings,rent,limiting_amount", + [ + (24_000, 36_000, "cap"), + (24_000, 7_300, "assistance"), + (40_000, 36_000, "zero"), + ], +) +def test_housing_cap_uses_computed_hud_benefits( + forecast, earnings, rent, limiting_amount +): + simulation = Simulation( + situation=household(county="06037", earnings=earnings, rent=rent), + spm={"geography_kind": "national"}, + ) + # Compute actual country formulas from earnings, rent and the PHA standard; + # neither hud_ttp nor housing_assistance is supplied as a test input. + tenant_payment = simulation.calculate("hud_ttp", YEAR).astype(np.float64) + assistance = simulation.calculate("housing_assistance", YEAR).astype(np.float64) + raw_housing = canonical(forecast)["housing_portion"] + raw_cap = np.maximum(raw_housing - tenant_payment, 0) + expected = np.minimum(assistance, raw_cap) + result = simulation.calculate("spm_unit_capped_housing_subsidy", YEAR) + np.testing.assert_array_equal(result, expected.astype(result.dtype)) + assert tenant_payment[0] > 0 + assert assistance[0] > 0 + if limiting_amount == "cap": + assert 0 < raw_cap[0] < assistance[0] + elif limiting_amount == "assistance": + assert 0 < assistance[0] < raw_cap[0] + else: + assert raw_cap[0] == 0 + + +@pytest.mark.parametrize( + "earnings,positive_cap", + [(26_085.168, True), (26_085.17, True), (26_085.172, False)], +) +def test_housing_cap_preserves_raw_amount_at_tenant_payment_boundary( + forecast, earnings, positive_cap +): + # Adjacent representable earnings put the computed HUD tenant payment just + # below, at, and above the stored 2024 national single-renter housing amount. + simulation = Simulation( + situation=household(county="06037", earnings=earnings), + spm={"geography_kind": "national"}, + ) + tenant_payment = simulation.calculate("hud_ttp", YEAR).astype(np.float64) + assistance = simulation.calculate("housing_assistance", YEAR).astype(np.float64) + stored_housing = simulation.calculate( + "spm_unit_spm_threshold_housing_portion", YEAR + ) + raw_housing = canonical(forecast)["housing_portion"] + expected = np.minimum(assistance, np.maximum(raw_housing - tenant_payment, 0)) + result = simulation.calculate("spm_unit_capped_housing_subsidy", YEAR) + np.testing.assert_array_equal(result, expected.astype(result.dtype)) + assert bool(result[0] > 0) is positive_cap + if positive_cap: + incorrectly_rounded = np.minimum( + assistance, + np.maximum(stored_housing.astype(np.float64) - tenant_payment, 0), + ).astype(result.dtype) + assert result[0] != incorrectly_rounded[0] + if earnings == 26_085.17: + assert stored_housing[0] == tenant_payment[0] + assert result[0] > 0 + + +def test_tax_only_state_input_is_lazy_and_spm_never_computes_county(monkeypatch): + simulation = Simulation(situation=household(earnings=50_000)) + assert simulation.calculate("income_tax", YEAR)[0] > 0 + assert simulation.spm_provenance()["years"] == {} + # Tax formulas may resolve their own county defaults. SPM must still use + # only county_fips input, even if an inferred county is already cached. + original_calculate = simulation.calculate + + def calculate(variable, *args, **kwargs): + assert variable not in {"county", "first_county_in_state"} + return original_calculate(variable, *args, **kwargs) + + monkeypatch.setattr(simulation, "calculate", calculate) + with pytest.raises(SPMInputError) as error: + simulation.calculate("spm_unit_spm_threshold", YEAR) + assert error.value.code == "SPM_GEOGRAPHY_REQUIRED" + assert simulation.calculate("income_tax", YEAR)[0] > 0 + + +@pytest.mark.parametrize("age", [14, 15, 17, 18]) +@pytest.mark.parametrize("role", [False, True]) +def test_spm_age_boundary_preserves_generic_demographics(age, role): + simulation = Simulation( + situation=household( + people={ + "person1": { + "age": {YEAR: age}, + "is_spm_independent_minor_role": role, + } + } + ), + spm={"geography_kind": "national"}, + ) + expected_spm_adult = age >= 18 or (age >= 15 and role) + assert simulation.calculate("spm_measurement_adults", YEAR)[0] == int( + expected_spm_adult + ) + assert simulation.calculate("spm_measurement_children", YEAR)[0] == int( + not expected_spm_adult + ) + assert bool(simulation.calculate("is_adult", YEAR)[0]) is (age >= 18) + assert bool(simulation.calculate("is_child", YEAR)[0]) is (age < 18) + assert simulation.calculate("spm_unit_count_adults", YEAR)[0] == int(age >= 18) + assert simulation.calculate("spm_unit_count_children", YEAR)[0] == int(age < 18) + if expected_spm_adult: + assert simulation.calculate("spm_unit_spm_threshold", YEAR)[0] > 0 + else: + with pytest.raises(SPMInputError) as error: + simulation.calculate("spm_unit_spm_threshold", YEAR) + assert error.value.to_dict()["code"] == "SPM_COMPOSITION_REQUIRED" + + +def test_household_roles_preserve_native_spm_membership(forecast): + people = { + "person1": {"age": {YEAR: 40}}, + "person2": {"age": {YEAR: 15}, "is_household_spouse": True}, + "person3": {"age": {YEAR: 17}, "is_household_head": True}, + "person4": {"age": {YEAR: 14}, "is_household_head": True}, + } + simulation = Simulation( + situation=household(people=people), spm={"geography_kind": "national"} + ) + np.testing.assert_array_equal( + simulation.calculate("is_spm_independent_minor_role", YEAR), + [False, True, True, True], + ) + assert simulation.calculate("spm_unit_size", YEAR)[0] == 4 + assert simulation.calculate("spm_unit_count_adults", YEAR)[0] == 1 + assert simulation.calculate("spm_unit_count_children", YEAR)[0] == 3 + assert simulation.calculate("spm_measurement_adults", YEAR)[0] == 3 + assert simulation.calculate("spm_measurement_children", YEAR)[0] == 1 + threshold = simulation.calculate("spm_unit_spm_threshold", YEAR) + expected = canonical(forecast, adults=3, children=1)["threshold"] + assert threshold[0] == np.asarray(expected, dtype=threshold.dtype) + + +def test_fixed_metro_matches_county_resolution_without_county_input(forecast): + assignment = forecast.resolve_county(YEAR, "06037") + assert assignment["kind"] == "metro" + config = {"geography_kind": "metro", "geography_id": assignment["area_id"]} + fixed = Simulation(situation=household(), spm=config) + observed = Simulation(situation=household(county="06037")) + fixed_threshold = fixed.calculate("spm_unit_spm_threshold", YEAR) + np.testing.assert_array_equal( + fixed_threshold, observed.calculate("spm_unit_spm_threshold", YEAR) + ) + expected = canonical(forecast, geography_kind="metro", area=assignment["area_id"])[ + "threshold" + ] + assert fixed_threshold[0] == np.asarray(expected, dtype=fixed_threshold.dtype) + assert fixed.spm_config["geography_id"] == assignment["area_id"] + provenance = json.loads(json.dumps(fixed.spm_provenance())) + assert provenance["geography_kind"] == "metro" + assert provenance["geographies"][0]["county_assignment"] is None + assert ( + observed.spm_provenance()["geographies"][0]["county_assignment"]["county_fips"] + == "06037" + ) + + +def test_county_location_requires_observed_input_instead_of_fixed_geography_id(): + with pytest.raises(ValueError, match="Only a fixed metro"): + Simulation( + situation=household(), + spm={"geography_kind": "county", "geography_id": "06037"}, + ) + + +@pytest.mark.parametrize("variable", ["household_net_income", "marginal_tax_rate"]) +@pytest.mark.parametrize("national", [False, True]) +def test_resource_consumers_without_housing_assistance_never_touch_geography( + variable, national +): + """A unit with nothing to cap must not need SPM geography or composition. + + Partners request net income, benefits and marginal rates without any SPM + measurement, and most households carry no housing assistance; the cap + consults the canonical housing portion only for units that do. + """ + situation = household(earnings=50_000) + situation["spm_units"]["spm_unit"]["receives_housing_assistance"] = {YEAR: False} + simulation = Simulation( + situation=situation, + spm={"geography_kind": "national"} if national else None, + ) + assert simulation.calculate("housing_assistance", YEAR)[0] == 0 + result = simulation.calculate(variable, YEAR) + assert np.all(np.isfinite(result)) + assert np.all(result > 0) + assert simulation.calculate("spm_unit_capped_housing_subsidy", YEAR)[0] == 0 + # No measurement was looked up, so nothing was received. + assert simulation.spm_provenance()["years"] == {} + + +@pytest.mark.parametrize("variable", ["household_net_income", "marginal_tax_rate"]) +def test_resource_consumers_with_housing_assistance_require_geography(variable): + """Once there is assistance to cap, the county requirement applies. + + Earnings stay low so the tenant payment sits below the housing portion + and the cap binds on the assistance rather than on zero. + """ + situation = household(earnings=6_000) + situation["spm_units"]["spm_unit"]["housing_assistance"] = {YEAR: 5_000} + with pytest.raises(SPMInputError) as error: + Simulation(situation=situation).calculate(variable, YEAR) + assert error.value.code == "SPM_GEOGRAPHY_REQUIRED" + assert error.value.to_dict()["code"] == "SPM_GEOGRAPHY_REQUIRED" + national = Simulation(situation=situation, spm={"geography_kind": "national"}) + result = national.calculate(variable, YEAR) + assert np.all(np.isfinite(result)) + capped = national.calculate("spm_unit_capped_housing_subsidy", YEAR)[0] + assert 0 < capped <= 5_000 + assert str(YEAR) in national.spm_provenance()["years"] + + +def two_household_population(*, assisted_county): + """Household A has housing assistance to cap; household B has none and no county.""" + return { + "people": { + "a": { + "age": {YEAR: 40}, + "employment_income": {YEAR: 6_000}, + "pre_subsidy_rent": {YEAR: 36_000}, + }, + "b": { + "age": {YEAR: 40}, + "employment_income": {YEAR: 50_000}, + "pre_subsidy_rent": {YEAR: 36_000}, + }, + }, + "households": { + "household_a": { + "members": ["a"], + "state_code": {YEAR: "CA"}, + "pha_payment_standard": {YEAR: 36_000}, + **( + {"county_fips": {YEAR: assisted_county}} + if assisted_county is not None + else {} + ), + }, + "household_b": { + "members": ["b"], + "state_code": {YEAR: "CA"}, + "pha_payment_standard": {YEAR: 36_000}, + }, + }, + "spm_units": { + "unit_a": { + "members": ["a"], + "spm_unit_tenure_type": {YEAR: "RENTER"}, + "housing_assistance": {YEAR: 5_000}, + }, + "unit_b": { + "members": ["b"], + "spm_unit_tenure_type": {YEAR: "RENTER"}, + "receives_housing_assistance": {YEAR: False}, + }, + }, + } + + +def test_housing_cap_evaluates_only_assisted_units_in_a_mixed_population(): + """Masked rows align: the assisted unit is capped, the other is zero and free. + + Unit B has no county and no assistance, so it must impose no requirement + on the population while unit A, which has both, is capped normally. + """ + simulation = Simulation(situation=two_household_population(assisted_county="06037")) + assert list(simulation.calculate("housing_assistance", YEAR)) == [5_000, 0] + capped = simulation.calculate("spm_unit_capped_housing_subsidy", YEAR) + assert capped.dtype == np.float32 + assert 0 < capped[0] <= 5_000 + assert capped[1] == 0 + net_income = simulation.calculate("household_net_income", YEAR) + assert np.all(np.isfinite(net_income)) and np.all(net_income > 0) + provenance = simulation.spm_provenance() + assert str(YEAR) in provenance["years"] + # Only the assisted unit's county reached the provider: one receipt, for + # unit A's Los Angeles County, resolved to its metropolitan area. + assert provenance["geography_kind"] == "county" + assert len(provenance["geographies"]) == 1 + (receipt,) = provenance["geographies"] + assert receipt["year"] == YEAR + assert receipt["tenure"] == "renter" + assert receipt["geography"]["area_id"] == "31080" + + +def test_assisted_unit_without_county_fails_closed_in_a_mixed_population(): + """An unassisted neighbour does not relax the requirement for the assisted unit.""" + simulation = Simulation(situation=two_household_population(assisted_county=None)) + with pytest.raises(SPMInputError) as error: + simulation.calculate("spm_unit_capped_housing_subsidy", YEAR) + assert error.value.code == "SPM_GEOGRAPHY_REQUIRED" + + +@pytest.mark.parametrize( + "state,variable,period,expected", + [ + # Colorado's need-standard table: zero adults, one child = $117/month. + ("CO", "co_tanf_need_standard", YEAR, 117 * 12), + # Minnesota MFIP excludes $100/month of child support for one child. + ("MN", "mn_mfip_child_support_income_exclusion", "2024-01", 100), + ], +) +def test_generic_benefit_consumers_do_not_use_spm_measurement_composition( + state, variable, period, expected +): + # Each unit has one 16-year-old. The second has a source-backed SPM role; + # the first has no classified SPM adult. Both still count as benefit children. + people = { + name: { + "age": {YEAR: 16}, + "is_spm_independent_minor_role": role, + "child_support_received": {YEAR: 1_800}, + } + for name, role in (("person1", False), ("person2", True)) + } + simulation = Simulation( + situation={ + "people": people, + "spm_units": {name: {"members": [name]} for name in people}, + "households": { + name: {"members": [name], "state_code": {YEAR: state}} + for name in people + }, + } + ) + np.testing.assert_array_equal( + simulation.calculate(variable, period), [expected, expected] + ) + np.testing.assert_array_equal( + simulation.calculate("spm_unit_count_adults", YEAR), [0, 0] + ) + np.testing.assert_array_equal( + simulation.calculate("spm_unit_count_children", YEAR), [1, 1] + ) + np.testing.assert_array_equal( + simulation.calculate("spm_measurement_adults", YEAR), [0, 1] + ) + np.testing.assert_array_equal( + simulation.calculate("spm_measurement_children", YEAR), [1, 0] + ) + assert simulation.spm_provenance()["years"] == {} diff --git a/policyengine_us/tools/default_uprating.py b/policyengine_us/tools/default_uprating.py index 731634fc6df..d26a56d6d5a 100644 --- a/policyengine_us/tools/default_uprating.py +++ b/policyengine_us/tools/default_uprating.py @@ -60,7 +60,6 @@ "casualty_loss", "partnership_s_corp_income", "taxable_interest_income", - "spm_unit_spm_threshold", "non_sch_d_capital_gains", "farm_operations_income", "taxable_403b_distributions", diff --git a/policyengine_us/variables/household/income/spm_unit/spm_unit_capped_housing_subsidy.py b/policyengine_us/variables/household/income/spm_unit/spm_unit_capped_housing_subsidy.py index cd411789750..b24f0ecf20b 100644 --- a/policyengine_us/variables/household/income/spm_unit/spm_unit_capped_housing_subsidy.py +++ b/policyengine_us/variables/household/income/spm_unit/spm_unit_capped_housing_subsidy.py @@ -1,4 +1,5 @@ from policyengine_us.model_api import * +from policyengine_us.spm import masked_policyengine_amount class spm_unit_capped_housing_subsidy(Variable): @@ -10,11 +11,20 @@ class spm_unit_capped_housing_subsidy(Variable): reference = "https://www2.census.gov/programs-surveys/supplemental-poverty-measure/datasets/spm/spm_techdoc.pdf" def formula(spm_unit, period, parameters): - housing_assistance = spm_unit("housing_assistance", period) - housing_portion = spm_unit( - "spm_unit_spm_threshold_housing_portion", - period, + housing_assistance = spm_unit("housing_assistance", period).astype("float64") + # Only a unit with housing assistance has anything to cap. Consulting + # the canonical housing portion for the other units would demand SPM + # geography and composition from resource and benefit consumers that + # never use the measurement, so the cap is evaluated for assisted units + # alone; every other unit's capped subsidy is zero by construction. + assisted = housing_assistance > 0 + if not assisted.any(): + return np.zeros_like(housing_assistance) + # Apply the country-owned cap to the unrounded canonical housing amount; + # the model stores this final benefit amount with one dtype conversion. + housing_portion = masked_policyengine_amount( + spm_unit, period, "housing_portion", assisted ) - tenant_payment = spm_unit("hud_ttp", period) + tenant_payment = spm_unit("hud_ttp", period).astype("float64") cap = max_(housing_portion - tenant_payment, 0) - return min_(housing_assistance, cap) + return where(assisted, min_(housing_assistance, cap), 0) diff --git a/policyengine_us/variables/household/income/spm_unit/spm_unit_geographic_adjustment.py b/policyengine_us/variables/household/income/spm_unit/spm_unit_geographic_adjustment.py deleted file mode 100644 index e99bbff635d..00000000000 --- a/policyengine_us/variables/household/income/spm_unit/spm_unit_geographic_adjustment.py +++ /dev/null @@ -1,46 +0,0 @@ -from policyengine_us.model_api import * -from policyengine_us.variables.household.income.spm_unit.spm_unit_tenure_type import ( - SPMUnitTenureType, -) -from spm_calculator.geoadj import get_cd_geoadj - - -SPM_TENURE_TYPE_TO_GEOADJ_KEY = { - SPMUnitTenureType.OWNER_WITH_MORTGAGE: "owner_with_mortgage", - SPMUnitTenureType.OWNER_WITHOUT_MORTGAGE: "owner_without_mortgage", - SPMUnitTenureType.RENTER: "renter", -} - - -def _cd_geoadj_array(cd_geoids, tenure, geoadj_year: int = 2023): - geoadj = np.ones(len(cd_geoids), dtype=float) - valid_cd = cd_geoids > 0 - - for tenure_enum, tenure_key in SPM_TENURE_TYPE_TO_GEOADJ_KEY.items(): - tenure_mask = valid_cd & (tenure == tenure_enum) - for cd_geoid in np.unique(cd_geoids[tenure_mask]): - try: - value = get_cd_geoadj( - int(cd_geoid), - year=geoadj_year, - tenure=tenure_key, - ) - except ValueError: - value = 1.0 - geoadj[tenure_mask & (cd_geoids == cd_geoid)] = value - - return geoadj - - -class spm_unit_geographic_adjustment(Variable): - value_type = float - entity = SPMUnit - label = "SPM unit geographic adjustment" - documentation = "Geographic adjustment applied to the SPM reference threshold." - definition_period = YEAR - default_value = 1.0 - - def formula_2015(spm_unit, period, parameters): - cd_geoids = spm_unit.household("congressional_district_geoid", period) - tenure = spm_unit("spm_unit_tenure_type", period) - return _cd_geoadj_array(cd_geoids, tenure) diff --git a/policyengine_us/variables/household/income/spm_unit/spm_unit_reference_spm_threshold.py b/policyengine_us/variables/household/income/spm_unit/spm_unit_reference_spm_threshold.py deleted file mode 100644 index 44b2366b0f5..00000000000 --- a/policyengine_us/variables/household/income/spm_unit/spm_unit_reference_spm_threshold.py +++ /dev/null @@ -1,67 +0,0 @@ -import numpy as np - -from policyengine_us.model_api import * -from policyengine_us.variables.household.income.spm_unit.spm_unit_tenure_type import ( - SPMUnitTenureType, -) -from spm_calculator.forecast import ( - HISTORICAL_THRESHOLDS, - get_latest_published_year, -) - - -LATEST_PUBLISHED_SPM_THRESHOLD_YEAR = get_latest_published_year() - - -def _reference_threshold_array(tenure, year: int, cpi_u_parameter): - """Published Betson reference thresholds for ``year``, uprated past - the latest BLS-published year using PolicyEngine's CPI-U parameter. - - The published values and three-tenure structure come from - ``spm-calculator`` so there is one source of truth across the - PolicyEngine stack. - """ - if year <= LATEST_PUBLISHED_SPM_THRESHOLD_YEAR: - thresholds = HISTORICAL_THRESHOLDS[year] - else: - factor = float( - cpi_u_parameter(f"{year}-02-01") - / cpi_u_parameter(f"{LATEST_PUBLISHED_SPM_THRESHOLD_YEAR}-02-01") - ) - thresholds = { - k: v * factor - for k, v in HISTORICAL_THRESHOLDS[ - LATEST_PUBLISHED_SPM_THRESHOLD_YEAR - ].items() - } - - values = np.full(len(tenure), thresholds["renter"], dtype=float) - values = np.where( - tenure == SPMUnitTenureType.OWNER_WITH_MORTGAGE, - thresholds["owner_with_mortgage"], - values, - ) - values = np.where( - tenure == SPMUnitTenureType.OWNER_WITHOUT_MORTGAGE, - thresholds["owner_without_mortgage"], - values, - ) - return values - - -class spm_unit_reference_spm_threshold(Variable): - value_type = float - entity = SPMUnit - label = "SPM unit reference SPM poverty threshold" - documentation = "SPM reference threshold by tenure, before equivalence-scale and geographic adjustments." - definition_period = YEAR - unit = USD - - def formula_2015(spm_unit, period, parameters): - cpi_u = parameters.gov.bls.cpi.cpi_u - tenure = spm_unit("spm_unit_tenure_type", period) - return _reference_threshold_array( - tenure, - period.start.year, - cpi_u, - ) diff --git a/policyengine_us/variables/household/income/spm_unit/spm_unit_spm_threshold.py b/policyengine_us/variables/household/income/spm_unit/spm_unit_spm_threshold.py deleted file mode 100644 index 47db6900831..00000000000 --- a/policyengine_us/variables/household/income/spm_unit/spm_unit_spm_threshold.py +++ /dev/null @@ -1,23 +0,0 @@ -from policyengine_us.model_api import * - - -class spm_unit_spm_threshold(Variable): - value_type = float - entity = SPMUnit - label = "SPM unit's SPM poverty threshold" - definition_period = YEAR - unit = USD - - def formula_2015(spm_unit, period, parameters): - """Rebuild the SPM threshold from current composition, current - tenure, and the unit-specific geographic adjustment. - - Base reference thresholds and the Betson three-parameter - equivalence scale come from ``spm-calculator``. - """ - unadjusted_threshold = spm_unit( - "spm_unit_unadjusted_spm_threshold", - period, - ) - geoadj = spm_unit("spm_unit_geographic_adjustment", period) - return unadjusted_threshold * geoadj diff --git a/policyengine_us/variables/household/income/spm_unit/spm_unit_spm_threshold_housing_portion.py b/policyengine_us/variables/household/income/spm_unit/spm_unit_spm_threshold_housing_portion.py deleted file mode 100644 index 259f0280df8..00000000000 --- a/policyengine_us/variables/household/income/spm_unit/spm_unit_spm_threshold_housing_portion.py +++ /dev/null @@ -1,52 +0,0 @@ -from policyengine_us.model_api import * -from policyengine_us.variables.household.income.spm_unit.spm_unit_tenure_type import ( - SPMUnitTenureType, -) -from spm_calculator.geoadj import get_housing_share - - -SPM_TENURE_TYPE_TO_HOUSING_SHARE = { - SPMUnitTenureType.OWNER_WITH_MORTGAGE: get_housing_share("owner_with_mortgage"), - SPMUnitTenureType.OWNER_WITHOUT_MORTGAGE: get_housing_share( - "owner_without_mortgage" - ), - SPMUnitTenureType.RENTER: get_housing_share("renter"), -} - - -def _housing_share_array(tenure): - values = np.full( - len(tenure), - SPM_TENURE_TYPE_TO_HOUSING_SHARE[SPMUnitTenureType.RENTER], - dtype=float, - ) - for tenure_type, housing_share in SPM_TENURE_TYPE_TO_HOUSING_SHARE.items(): - values = np.where(tenure == tenure_type, housing_share, values) - return values - - -class spm_unit_spm_threshold_housing_portion(Variable): - value_type = float - entity = SPMUnit - label = "SPM unit SPM threshold housing portion" - documentation = ( - "Geographically adjusted housing portion of the SPM poverty threshold." - ) - definition_period = YEAR - unit = USD - - def formula_2015(spm_unit, period, parameters): - unadjusted_threshold = spm_unit( - "spm_unit_unadjusted_spm_threshold", - period, - ) - geoadj = spm_unit("spm_unit_geographic_adjustment", period) - tenure = spm_unit("spm_unit_tenure_type", period) - housing_share = _housing_share_array(tenure) - - # The threshold applies geographic adjustment only to the housing - # share: adjusted = unadjusted * (1 - h + h * local_rent_ratio). - # Since geoadj = 1 - h + h * local_rent_ratio, the adjusted housing - # share is h * local_rent_ratio = geoadj + h - 1. - geoadjusted_housing_share = geoadj + housing_share - 1 - return unadjusted_threshold * geoadjusted_housing_share diff --git a/policyengine_us/variables/household/income/spm_unit/spm_unit_unadjusted_spm_threshold.py b/policyengine_us/variables/household/income/spm_unit/spm_unit_unadjusted_spm_threshold.py deleted file mode 100644 index 07e7e354fa6..00000000000 --- a/policyengine_us/variables/household/income/spm_unit/spm_unit_unadjusted_spm_threshold.py +++ /dev/null @@ -1,21 +0,0 @@ -from policyengine_us.model_api import * -from spm_calculator.equivalence_scale import spm_equivalence_scale - - -class spm_unit_unadjusted_spm_threshold(Variable): - value_type = float - entity = SPMUnit - label = "SPM unit unadjusted SPM poverty threshold" - documentation = "SPM poverty threshold before geographic adjustment." - definition_period = YEAR - unit = USD - - def formula_2015(spm_unit, period, parameters): - reference_threshold = spm_unit( - "spm_unit_reference_spm_threshold", - period, - ) - adults = spm_unit("spm_unit_count_adults", period) - children = spm_unit("spm_unit_count_children", period) - equivalence_scale = spm_equivalence_scale(adults, children) - return reference_threshold * equivalence_scale diff --git a/pyproject.toml b/pyproject.toml index ffb1ad0309a..320fd79b852 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -24,7 +24,7 @@ dependencies = [ # extension and is available on every supported Python version. "pandas>=3.0", "policyengine-core>=3.30.1", - "spm-calculator==0.3.1", + "spm-calculator==1.0.0", "tables>=3.9", "tqdm>=4.67.1", ] diff --git a/scripts/benchmark_household_simulation.py b/scripts/benchmark_household_simulation.py new file mode 100644 index 00000000000..7b3ba3729d4 --- /dev/null +++ b/scripts/benchmark_household_simulation.py @@ -0,0 +1,101 @@ +"""Time repeated single-household simulations of the shipped policy. + +Run the same file against two checkouts to compare them: + + uv run --locked python scripts/benchmark_household_simulation.py + +Each iteration constructs one ``Simulation`` from an identical situation and +calculates ``household_net_income``, which is what a household API request +does. The situation carries a county so the run is identical on revisions that +require an explicit SPM area and on revisions that ignore county entirely. +""" + +import argparse +import json +import time + + +YEAR = 2024 +COUNTY = "06037" + + +def situation(earnings): + return { + "people": { + "you": { + "age": {YEAR: 40}, + "employment_income": {YEAR: earnings}, + }, + "spouse": {"age": {YEAR: 38}}, + "child": {"age": {YEAR: 6}}, + }, + "families": {"family": {"members": ["you", "spouse", "child"]}}, + "marital_units": {"marital_unit": {"members": ["you", "spouse"]}}, + "tax_units": {"tax_unit": {"members": ["you", "spouse", "child"]}}, + "spm_units": {"spm_unit": {"members": ["you", "spouse", "child"]}}, + "households": { + "household": { + "members": ["you", "spouse", "child"], + "state_code": {YEAR: "CA"}, + "county_fips": {YEAR: COUNTY}, + } + }, + } + + +def main(): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--iterations", type=int, default=20) + arguments = parser.parse_args() + + import_start = time.perf_counter() + from policyengine_us import Simulation + + import_seconds = time.perf_counter() - import_start + + iterations = [] + results = [] + for index in range(arguments.iterations): + # Vary earnings so no iteration can reuse another's cached result. + start = time.perf_counter() + simulation = Simulation(situation=situation(50_000 + 1_000 * index)) + constructed = time.perf_counter() + net_income = float(simulation.calculate("household_net_income", YEAR)[0]) + done = time.perf_counter() + iterations.append( + { + "construct_seconds": constructed - start, + "calculate_seconds": done - constructed, + "total_seconds": done - start, + } + ) + results.append(net_income) + + totals = [iteration["total_seconds"] for iteration in iterations] + print( + json.dumps( + { + "iterations": len(iterations), + "import_seconds": round(import_seconds, 3), + "first_iteration_seconds": round(totals[0], 3), + "remaining_mean_seconds": round( + sum(totals[1:]) / max(len(totals) - 1, 1), 3 + ), + "total_seconds": round(sum(totals), 3), + "construct_total_seconds": round( + sum(i["construct_seconds"] for i in iterations), 3 + ), + "calculate_total_seconds": round( + sum(i["calculate_seconds"] for i in iterations), 3 + ), + "net_income_first": results[0], + "net_income_last": results[-1], + "per_iteration_seconds": [round(total, 3) for total in totals], + }, + indent=2, + ) + ) + + +if __name__ == "__main__": + main() diff --git a/uv.lock b/uv.lock index 16839123233..1f19f25b1a0 100644 --- a/uv.lock +++ b/uv.lock @@ -1730,7 +1730,7 @@ wheels = [ [[package]] name = "policyengine-us" -version = "1.824.8" +version = "1.825.2" source = { editable = "." } dependencies = [ { name = "microdf-python" }, @@ -1763,7 +1763,7 @@ requires-dist = [ { name = "policyengine-core", specifier = ">=3.30.1" }, { name = "ruff", marker = "extra == 'dev'", specifier = ">=0.9.0" }, { name = "setuptools", marker = "extra == 'dev'", specifier = ">=80.9.0" }, - { name = "spm-calculator", specifier = "==0.3.1" }, + { name = "spm-calculator", specifier = "==1.0.0" }, { name = "tables", specifier = ">=3.9" }, { name = "towncrier", marker = "extra == 'dev'", specifier = ">=24.8.0" }, { name = "tqdm", specifier = ">=4.67.1" }, @@ -2443,7 +2443,7 @@ wheels = [ [[package]] name = "spm-calculator" -version = "0.3.1" +version = "1.0.0" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "census" }, @@ -2453,9 +2453,9 @@ dependencies = [ { name = "requests" }, { name = "us" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/54/3b/b805c7e3e18c5b5c00f61b60112f9690d084c910e2481bc020f35390d8fd/spm_calculator-0.3.1.tar.gz", hash = "sha256:41f2f4d00d8c03422a7d57b800052e7760b88e463a5884802f83ed58d35c18c1", size = 75945, upload-time = "2026-04-17T19:52:39.707Z" } +sdist = { url = "https://files.pythonhosted.org/packages/c2/5e/3f28b212c401795990250bb3daec741f699a8c7dc2a23a3db7005df6feb6/spm_calculator-1.0.0.tar.gz", hash = "sha256:a99aac8c2c0bf81a9455105bbf872366f1ea9066a99cd793714881cd3db846dc", size = 7150825, upload-time = "2026-09-11T15:45:12.732Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/8e/1b/29f705f8a96fc7f55f2c07dfcddbbae78efdc6f174d25d4a0560fc3f5cf9/spm_calculator-0.3.1-py3-none-any.whl", hash = "sha256:52c57ecc5a240ec941b0f2b0d93bc4fa437ef6250e233baed8e11916fa9c1150", size = 57826, upload-time = "2026-04-17T19:52:38.444Z" }, + { url = "https://files.pythonhosted.org/packages/19/a0/c484f69a0ebf88a9b9fd0ac28176f8df66550f46714f1b0c5ec0b360824e/spm_calculator-1.0.0-py3-none-any.whl", hash = "sha256:e354937a5e1a4045d4966ed594a528d8b02866fabaac9bb5672017004b627305", size = 7395384, upload-time = "2026-09-11T15:45:10.492Z" }, ] [[package]]