From c52f4b1c4afa28c499bdcf84b725e32b338dd530 Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Sat, 18 Apr 2026 22:08:02 -0400 Subject: [PATCH] Move provenance out of core; lazy-import h5py Separates release-manifest + TRACE TRO emission from the core value layer. Consumers who only need Simulation / Policy / Variable / Parameter no longer transitively import h5py through scoping_strategy / constituency_impact / local_authority_impact. File moves: - core/release_manifest.py -> provenance/manifest.py - core/trace_tro.py -> provenance/trace.py New provenance/__init__.py re-exports the public surface (get_release_manifest, build_trace_tro_from_release_bundle, serialize_trace_tro, canonical_json_bytes, etc.). core/__init__.py drops the 20 provenance re-exports and keeps only value objects (Dataset, Variable, Parameter*, Policy, Dynamic, Simulation, Region, scoping strategies, TaxBenefitModel, TaxBenefitModelVersion). Explicit core -> provenance import in tax_benefit_model_version.py. Lazy h5py: - core/scoping_strategy.py: h5py no longer at top of module; imported inside WeightReplacementStrategy.apply() only. - outputs/constituency_impact.py: same. - outputs/local_authority_impact.py: same. Internal callers migrated: - tax_benefit_models/{us,uk}/model.py - tax_benefit_models/{us,uk}/datasets.py - countries/{us,uk}/regions.py - cli.py - results/trace_tro.py - scripts/generate_trace_tros.py - tests/test_{release_manifests,trace_tro,manifest_version_mismatch}.py - docs/release-bundles.md 216 tests pass locally across the v4 surface. `from policyengine.core import Simulation` + `from policyengine.provenance import get_release_manifest` both work without h5py installed (verified by temporarily uninstalling and retrying). The full `import policyengine as pe` still pulls h5py because policyengine_us / policyengine_uk import it eagerly (upstream); that's outside our control. Co-Authored-By: Claude Opus 4.7 (1M context) --- changelog.d/v4-provenance-package.changed.md | 24 +++++ docs/release-bundles.md | 6 +- scripts/generate_trace_tros.py | 4 +- src/policyengine/cli.py | 6 +- src/policyengine/core/__init__.py | 37 ++------ src/policyengine/core/scoping_strategy.py | 7 +- .../core/tax_benefit_model_version.py | 5 +- src/policyengine/countries/uk/regions.py | 2 +- src/policyengine/countries/us/regions.py | 2 +- .../outputs/constituency_impact.py | 3 +- .../outputs/local_authority_impact.py | 3 +- src/policyengine/provenance/__init__.py | 89 +++++++++++++++++++ .../manifest.py} | 0 .../trace_tro.py => provenance/trace.py} | 2 +- src/policyengine/results/trace_tro.py | 4 +- .../tax_benefit_models/uk/datasets.py | 2 +- .../tax_benefit_models/uk/model.py | 2 +- .../tax_benefit_models/us/datasets.py | 2 +- .../tax_benefit_models/us/model.py | 2 +- tests/test_manifest_version_mismatch.py | 2 +- tests/test_release_manifests.py | 24 ++--- tests/test_trace_tro.py | 16 ++-- 22 files changed, 171 insertions(+), 73 deletions(-) create mode 100644 changelog.d/v4-provenance-package.changed.md create mode 100644 src/policyengine/provenance/__init__.py rename src/policyengine/{core/release_manifest.py => provenance/manifest.py} (100%) rename src/policyengine/{core/trace_tro.py => provenance/trace.py} (99%) diff --git a/changelog.d/v4-provenance-package.changed.md b/changelog.d/v4-provenance-package.changed.md new file mode 100644 index 00000000..8c016e02 --- /dev/null +++ b/changelog.d/v4-provenance-package.changed.md @@ -0,0 +1,24 @@ +**BREAKING (v4):** Separate the provenance layer from the core +value-object layer. + +- ``policyengine/core/release_manifest.py`` → ``policyengine/provenance/manifest.py`` +- ``policyengine/core/trace_tro.py`` → ``policyengine/provenance/trace.py`` +- New ``policyengine.provenance`` package re-exports the public + surface (``get_release_manifest``, ``get_data_release_manifest``, + ``build_trace_tro_from_release_bundle``, ``build_simulation_trace_tro``, + ``serialize_trace_tro``, ``canonical_json_bytes``, + ``compute_trace_composition_fingerprint``, etc.). +- ``policyengine.core`` no longer re-exports provenance types. + ``policyengine.core`` shrinks to value objects only (Dataset, + Variable, Parameter, Policy, Dynamic, Simulation, Region, + TaxBenefitModel, TaxBenefitModelVersion, scoping strategies). +- ``import policyengine.core.scoping_strategy`` no longer imports + ``h5py`` at module load; the weight-replacement code path + lazy-imports it. ``import policyengine.outputs.constituency_impact`` + and ``import policyengine.outputs.local_authority_impact`` do the + same. +- Migration for downstream: replace + ``from policyengine.core import DataReleaseManifest`` (et al.) + with ``from policyengine.provenance import DataReleaseManifest``. + The country-module imports in internal code (``tax_benefit_models/{us,uk}/model.py`` + and ``datasets.py``) are already updated. diff --git a/docs/release-bundles.md b/docs/release-bundles.md index ea014c9d..a28e0d1d 100644 --- a/docs/release-bundles.md +++ b/docs/release-bundles.md @@ -224,8 +224,8 @@ That is a country-data concern and lives in those repos. From Python: ```python -from policyengine.core.release_manifest import get_data_release_manifest, get_release_manifest -from policyengine.core.trace_tro import build_trace_tro_from_release_bundle, serialize_trace_tro +from policyengine.provenance.manifest import get_data_release_manifest, get_release_manifest +from policyengine.provenance.trace import build_trace_tro_from_release_bundle, serialize_trace_tro country = get_release_manifest("us") tro = build_trace_tro_from_release_bundle(country, get_data_release_manifest("us")) @@ -286,7 +286,7 @@ should run): ```python import hashlib, json, requests -from policyengine.core.trace_tro import canonical_json_bytes +from policyengine.provenance.trace import canonical_json_bytes sim_tro = json.load(open("results.trace.tro.jsonld")) perf = sim_tro["@graph"][0]["trov:hasPerformance"] diff --git a/scripts/generate_trace_tros.py b/scripts/generate_trace_tros.py index dce7ae8e..f9533bd9 100644 --- a/scripts/generate_trace_tros.py +++ b/scripts/generate_trace_tros.py @@ -18,12 +18,12 @@ import sys from pathlib import Path -from policyengine.core.release_manifest import ( +from policyengine.provenance.manifest import ( DataReleaseManifestUnavailableError, get_data_release_manifest, get_release_manifest, ) -from policyengine.core.trace_tro import ( +from policyengine.provenance.trace import ( build_trace_tro_from_release_bundle, serialize_trace_tro, ) diff --git a/src/policyengine/cli.py b/src/policyengine/cli.py index add36388..3a659643 100644 --- a/src/policyengine/cli.py +++ b/src/policyengine/cli.py @@ -6,7 +6,7 @@ - ``trace-tro-validate `` validate a TRO against the shipped schema - ``release-manifest `` print the bundled country manifest -See :mod:`policyengine.core.trace_tro` and ``docs/release-bundles.md``. +See :mod:`policyengine.provenance.trace` and ``docs/release-bundles.md``. """ from __future__ import annotations @@ -18,11 +18,11 @@ from pathlib import Path from typing import Optional, Sequence -from policyengine.core.release_manifest import ( +from policyengine.provenance.manifest import ( get_data_release_manifest, get_release_manifest, ) -from policyengine.core.trace_tro import ( +from policyengine.provenance.trace import ( build_trace_tro_from_release_bundle, serialize_trace_tro, ) diff --git a/src/policyengine/core/__init__.py b/src/policyengine/core/__init__.py index 71ca0132..4f749de4 100644 --- a/src/policyengine/core/__init__.py +++ b/src/policyengine/core/__init__.py @@ -1,3 +1,11 @@ +"""Core value objects: Dataset, Variable, Parameter, Policy, Simulation, Region. + +Provenance (release manifests, TRACE TROs) lives in +:mod:`policyengine.provenance` and is intentionally not re-exported +here — importing a core value object should not pull in the +provenance layer. +""" + from .dataset import Dataset from .dataset import YearData as YearData from .dataset import map_to_entity as map_to_entity @@ -11,22 +19,6 @@ from .region import Region as Region from .region import RegionRegistry as RegionRegistry from .region import RegionType as RegionType -from .release_manifest import CertifiedDataArtifact as CertifiedDataArtifact -from .release_manifest import CountryReleaseManifest as CountryReleaseManifest -from .release_manifest import DataBuildInfo as DataBuildInfo -from .release_manifest import DataCertification as DataCertification -from .release_manifest import DataPackageVersion as DataPackageVersion -from .release_manifest import DataReleaseArtifact as DataReleaseArtifact -from .release_manifest import DataReleaseManifest as DataReleaseManifest -from .release_manifest import PackageVersion as PackageVersion -from .release_manifest import ( - certify_data_release_compatibility as certify_data_release_compatibility, -) -from .release_manifest import get_data_release_manifest as get_data_release_manifest -from .release_manifest import get_release_manifest as get_release_manifest -from .release_manifest import ( - resolve_managed_dataset_reference as resolve_managed_dataset_reference, -) from .scoping_strategy import RegionScopingStrategy as RegionScopingStrategy from .scoping_strategy import RowFilterStrategy as RowFilterStrategy from .scoping_strategy import ScopingStrategy as ScopingStrategy @@ -38,19 +30,6 @@ from .tax_benefit_model_version import ( TaxBenefitModelVersion as TaxBenefitModelVersion, ) -from .trace_tro import ( - build_simulation_trace_tro as build_simulation_trace_tro, -) -from .trace_tro import ( - build_trace_tro_from_release_bundle as build_trace_tro_from_release_bundle, -) -from .trace_tro import ( - compute_trace_composition_fingerprint as compute_trace_composition_fingerprint, -) -from .trace_tro import ( - extract_bundle_tro_reference as extract_bundle_tro_reference, -) -from .trace_tro import serialize_trace_tro as serialize_trace_tro from .variable import Variable as Variable # Rebuild models to resolve forward references diff --git a/src/policyengine/core/scoping_strategy.py b/src/policyengine/core/scoping_strategy.py index 7d9b5126..96bbf409 100644 --- a/src/policyengine/core/scoping_strategy.py +++ b/src/policyengine/core/scoping_strategy.py @@ -14,7 +14,6 @@ from pathlib import Path from typing import Annotated, Literal, Optional, Union -import h5py import numpy as np import pandas as pd from microdf import MicroDataFrame @@ -127,7 +126,11 @@ def apply( region_id = self._find_region_index(lookup_df, self.region_code) - # Download weight matrix and extract weights for this region + # Download weight matrix and extract weights for this region. + # h5py is only needed here, so import lazily to keep + # `from policyengine.core import ...` light. + import h5py + weights_path = download_gcs_file( bucket=self.weight_matrix_bucket, file_path=self.weight_matrix_key, diff --git a/src/policyengine/core/tax_benefit_model_version.py b/src/policyengine/core/tax_benefit_model_version.py index eeddef85..5eb8f525 100644 --- a/src/policyengine/core/tax_benefit_model_version.py +++ b/src/policyengine/core/tax_benefit_model_version.py @@ -4,14 +4,15 @@ from pydantic import BaseModel, Field -from .release_manifest import ( +from policyengine.provenance.manifest import ( CountryReleaseManifest, DataCertification, PackageVersion, get_data_release_manifest, ) +from policyengine.provenance.trace import build_trace_tro_from_release_bundle + from .tax_benefit_model import TaxBenefitModel -from .trace_tro import build_trace_tro_from_release_bundle if TYPE_CHECKING: from .parameter import Parameter diff --git a/src/policyengine/countries/uk/regions.py b/src/policyengine/countries/uk/regions.py index d90f0ad0..32430d48 100644 --- a/src/policyengine/countries/uk/regions.py +++ b/src/policyengine/countries/uk/regions.py @@ -15,11 +15,11 @@ from typing import TYPE_CHECKING from policyengine.core.region import Region, RegionRegistry -from policyengine.core.release_manifest import resolve_region_dataset_path from policyengine.core.scoping_strategy import ( RowFilterStrategy, WeightReplacementStrategy, ) +from policyengine.provenance.manifest import resolve_region_dataset_path if TYPE_CHECKING: pass diff --git a/src/policyengine/countries/us/regions.py b/src/policyengine/countries/us/regions.py index 9e20d8b3..ca2f6b4f 100644 --- a/src/policyengine/countries/us/regions.py +++ b/src/policyengine/countries/us/regions.py @@ -8,8 +8,8 @@ """ from policyengine.core.region import Region, RegionRegistry -from policyengine.core.release_manifest import resolve_region_dataset_path from policyengine.core.scoping_strategy import RowFilterStrategy +from policyengine.provenance.manifest import resolve_region_dataset_path from .data import AT_LARGE_STATES, DISTRICT_COUNTS, US_PLACES, US_STATES diff --git a/src/policyengine/outputs/constituency_impact.py b/src/policyengine/outputs/constituency_impact.py index 60f76e0b..02e1bdfd 100644 --- a/src/policyengine/outputs/constituency_impact.py +++ b/src/policyengine/outputs/constituency_impact.py @@ -7,7 +7,6 @@ from typing import TYPE_CHECKING, Optional -import h5py import numpy as np import pandas as pd from pydantic import ConfigDict @@ -43,6 +42,8 @@ def run(self) -> None: constituency_df = pd.read_csv(self.constituency_csv_path) # Load weight matrix: shape (N_constituencies, N_households) + import h5py + with h5py.File(self.weight_matrix_path, "r") as f: weight_matrix = f[self.year][...] diff --git a/src/policyengine/outputs/local_authority_impact.py b/src/policyengine/outputs/local_authority_impact.py index 20b17efe..a4850dbf 100644 --- a/src/policyengine/outputs/local_authority_impact.py +++ b/src/policyengine/outputs/local_authority_impact.py @@ -7,7 +7,6 @@ from typing import TYPE_CHECKING, Optional -import h5py import numpy as np import pandas as pd from pydantic import ConfigDict @@ -43,6 +42,8 @@ def run(self) -> None: la_df = pd.read_csv(self.local_authority_csv_path) # Load weight matrix: shape (N_local_authorities, N_households) + import h5py + with h5py.File(self.weight_matrix_path, "r") as f: weight_matrix = f[self.year][...] diff --git a/src/policyengine/provenance/__init__.py b/src/policyengine/provenance/__init__.py new file mode 100644 index 00000000..548b7fc6 --- /dev/null +++ b/src/policyengine/provenance/__init__.py @@ -0,0 +1,89 @@ +"""Release-bundle provenance + TRACE TRO emission. + +Separated from :mod:`policyengine.core` so the value-object layer +(Dataset, Variable, Parameter, Policy, Simulation, Region) doesn't +force provenance imports on every consumer. + +.. code-block:: python + + from policyengine.provenance import ( + get_release_manifest, + get_data_release_manifest, + build_trace_tro_from_release_bundle, + build_simulation_trace_tro, + serialize_trace_tro, + ) +""" + +from .manifest import ( + CertifiedDataArtifact as CertifiedDataArtifact, +) +from .manifest import ( + CountryReleaseManifest as CountryReleaseManifest, +) +from .manifest import ( + DataBuildInfo as DataBuildInfo, +) +from .manifest import ( + DataCertification as DataCertification, +) +from .manifest import ( + DataPackageVersion as DataPackageVersion, +) +from .manifest import ( + DataReleaseArtifact as DataReleaseArtifact, +) +from .manifest import ( + DataReleaseManifest as DataReleaseManifest, +) +from .manifest import ( + DataReleaseManifestUnavailableError as DataReleaseManifestUnavailableError, +) +from .manifest import ( + PackageVersion as PackageVersion, +) +from .manifest import ( + certify_data_release_compatibility as certify_data_release_compatibility, +) +from .manifest import ( + fetch_pypi_wheel_metadata as fetch_pypi_wheel_metadata, +) +from .manifest import ( + get_data_release_manifest as get_data_release_manifest, +) +from .manifest import ( + get_release_manifest as get_release_manifest, +) +from .manifest import ( + https_dataset_uri as https_dataset_uri, +) +from .manifest import ( + https_release_manifest_uri as https_release_manifest_uri, +) +from .manifest import ( + resolve_dataset_reference as resolve_dataset_reference, +) +from .manifest import ( + resolve_local_managed_dataset_source as resolve_local_managed_dataset_source, +) +from .manifest import ( + resolve_managed_dataset_reference as resolve_managed_dataset_reference, +) +from .trace import ( + build_simulation_trace_tro as build_simulation_trace_tro, +) +from .trace import ( + build_trace_tro_from_release_bundle as build_trace_tro_from_release_bundle, +) +from .trace import ( + canonical_json_bytes as canonical_json_bytes, +) +from .trace import ( + compute_trace_composition_fingerprint as compute_trace_composition_fingerprint, +) +from .trace import ( + extract_bundle_tro_reference as extract_bundle_tro_reference, +) +from .trace import ( + serialize_trace_tro as serialize_trace_tro, +) diff --git a/src/policyengine/core/release_manifest.py b/src/policyengine/provenance/manifest.py similarity index 100% rename from src/policyengine/core/release_manifest.py rename to src/policyengine/provenance/manifest.py diff --git a/src/policyengine/core/trace_tro.py b/src/policyengine/provenance/trace.py similarity index 99% rename from src/policyengine/core/trace_tro.py rename to src/policyengine/provenance/trace.py index 76f1661a..83ac0b5b 100644 --- a/src/policyengine/core/trace_tro.py +++ b/src/policyengine/provenance/trace.py @@ -24,7 +24,7 @@ from collections.abc import Iterable, Mapping from typing import Any, Optional -from .release_manifest import ( +from .manifest import ( CountryReleaseManifest, DataCertification, DataReleaseManifest, diff --git a/src/policyengine/results/trace_tro.py b/src/policyengine/results/trace_tro.py index d904d5b4..85c7aed8 100644 --- a/src/policyengine/results/trace_tro.py +++ b/src/policyengine/results/trace_tro.py @@ -5,7 +5,7 @@ specific reform + ``results.json`` payload so a published result can be cited with an immutable composition fingerprint. -See :mod:`policyengine.core.trace_tro` for the bundle-level layer. +See :mod:`policyengine.provenance.trace` for the bundle-level layer. """ from __future__ import annotations @@ -14,7 +14,7 @@ from pathlib import Path from typing import Optional, Union -from policyengine.core.trace_tro import ( +from policyengine.provenance.trace import ( build_simulation_trace_tro, serialize_trace_tro, ) diff --git a/src/policyengine/tax_benefit_models/uk/datasets.py b/src/policyengine/tax_benefit_models/uk/datasets.py index 47f78403..e7207da7 100644 --- a/src/policyengine/tax_benefit_models/uk/datasets.py +++ b/src/policyengine/tax_benefit_models/uk/datasets.py @@ -6,7 +6,7 @@ from pydantic import ConfigDict from policyengine.core import Dataset, YearData -from policyengine.core.release_manifest import ( +from policyengine.provenance.manifest import ( dataset_logical_name, resolve_dataset_reference, ) diff --git a/src/policyengine/tax_benefit_models/uk/model.py b/src/policyengine/tax_benefit_models/uk/model.py index ce6f2dd9..c03ed90d 100644 --- a/src/policyengine/tax_benefit_models/uk/model.py +++ b/src/policyengine/tax_benefit_models/uk/model.py @@ -14,7 +14,7 @@ TaxBenefitModelVersion, Variable, ) -from policyengine.core.release_manifest import ( +from policyengine.provenance.manifest import ( certify_data_release_compatibility, dataset_logical_name, get_release_manifest, diff --git a/src/policyengine/tax_benefit_models/us/datasets.py b/src/policyengine/tax_benefit_models/us/datasets.py index da10733b..014309db 100644 --- a/src/policyengine/tax_benefit_models/us/datasets.py +++ b/src/policyengine/tax_benefit_models/us/datasets.py @@ -7,7 +7,7 @@ from pydantic import ConfigDict from policyengine.core import Dataset, YearData -from policyengine.core.release_manifest import ( +from policyengine.provenance.manifest import ( dataset_logical_name, resolve_dataset_reference, ) diff --git a/src/policyengine/tax_benefit_models/us/model.py b/src/policyengine/tax_benefit_models/us/model.py index cd56df09..ec3dd9e6 100644 --- a/src/policyengine/tax_benefit_models/us/model.py +++ b/src/policyengine/tax_benefit_models/us/model.py @@ -14,7 +14,7 @@ TaxBenefitModelVersion, Variable, ) -from policyengine.core.release_manifest import ( +from policyengine.provenance.manifest import ( certify_data_release_compatibility, dataset_logical_name, get_release_manifest, diff --git a/tests/test_manifest_version_mismatch.py b/tests/test_manifest_version_mismatch.py index f9145556..1c65230e 100644 --- a/tests/test_manifest_version_mismatch.py +++ b/tests/test_manifest_version_mismatch.py @@ -26,7 +26,7 @@ import warnings from unittest.mock import patch -from policyengine.core.release_manifest import get_release_manifest +from policyengine.provenance.manifest import get_release_manifest def _pick_mismatched_version(manifest_version: str) -> str: diff --git a/tests/test_release_manifests.py b/tests/test_release_manifests.py index 18d6eed3..1a07401b 100644 --- a/tests/test_release_manifests.py +++ b/tests/test_release_manifests.py @@ -5,7 +5,9 @@ from requests import Timeout -from policyengine.core.release_manifest import ( +from policyengine.core.tax_benefit_model import TaxBenefitModel +from policyengine.core.tax_benefit_model_version import TaxBenefitModelVersion +from policyengine.provenance.manifest import ( DataCertification, DataReleaseManifestUnavailableError, certify_data_release_compatibility, @@ -15,8 +17,6 @@ resolve_dataset_reference, resolve_managed_dataset_reference, ) -from policyengine.core.tax_benefit_model import TaxBenefitModel -from policyengine.core.tax_benefit_model_version import TaxBenefitModelVersion from policyengine.tax_benefit_models.uk import ( managed_microsimulation as managed_uk_microsimulation, ) @@ -179,7 +179,7 @@ def test__given_country__then_can_fetch_data_release_manifest(self): } with patch( - "policyengine.core.release_manifest.requests.get", + "policyengine.provenance.manifest.requests.get", return_value=_response_with_json(payload), ) as mock_get: manifest = get_data_release_manifest("us") @@ -204,7 +204,7 @@ def test__given_missing_data_release_manifest__then_fetch_raises_unavailable(sel response.status_code = 404 with patch( - "policyengine.core.release_manifest.requests.get", + "policyengine.provenance.manifest.requests.get", return_value=response, ): try: @@ -243,7 +243,7 @@ def test__given_range_specifier__then_certification_accepts_compatible_version( } with patch( - "policyengine.core.release_manifest.requests.get", + "policyengine.provenance.manifest.requests.get", return_value=_response_with_json(payload), ): certification = certify_data_release_compatibility( @@ -277,7 +277,7 @@ def test__given_matching_fingerprint__then_certification_allows_reuse(self): } with patch( - "policyengine.core.release_manifest.requests.get", + "policyengine.provenance.manifest.requests.get", return_value=_response_with_json(payload), ): certification = certify_data_release_compatibility( @@ -297,7 +297,7 @@ def test__given_private_manifest_unavailable__then_bundled_certification_is_used get_data_release_manifest.cache_clear() with patch( - "policyengine.core.release_manifest.get_data_release_manifest", + "policyengine.provenance.manifest.get_data_release_manifest", side_effect=DataReleaseManifestUnavailableError("private repo"), ): certification = certify_data_release_compatibility( @@ -314,11 +314,11 @@ def test__given_private_manifest_unavailable_and_fingerprint_mismatch__then_fail with ( patch( - "policyengine.core.release_manifest.get_data_release_manifest", + "policyengine.provenance.manifest.get_data_release_manifest", side_effect=DataReleaseManifestUnavailableError("private repo"), ), patch( - "policyengine.core.release_manifest.get_release_manifest", + "policyengine.provenance.manifest.get_release_manifest", return_value=MagicMock( certification=DataCertification( compatibility_basis="matching_data_build_fingerprint", @@ -345,7 +345,7 @@ def test__given_manifest_fetch_failure__then_certification_does_not_fallback( get_data_release_manifest.cache_clear() with patch( - "policyengine.core.release_manifest.get_data_release_manifest", + "policyengine.provenance.manifest.get_data_release_manifest", side_effect=Timeout("network timeout"), ): try: @@ -381,7 +381,7 @@ def test__given_mismatched_version_and_fingerprint__then_certification_fails(sel } with patch( - "policyengine.core.release_manifest.requests.get", + "policyengine.provenance.manifest.requests.get", return_value=_response_with_json(payload), ): try: diff --git a/tests/test_trace_tro.py b/tests/test_trace_tro.py index f78b4f33..9f32817f 100644 --- a/tests/test_trace_tro.py +++ b/tests/test_trace_tro.py @@ -1,6 +1,6 @@ """Tests for TRACE Transparent Research Object (TRO) export. -Covers bundle-level TROs (``policyengine.core.trace_tro``) and per-simulation +Covers bundle-level TROs (``policyengine.provenance.trace``) and per-simulation TROs (``policyengine.results.trace_tro``), plus the ``policyengine trace-tro`` CLI, determinism guarantees, and JSON-Schema conformance against TROv 2023/05. """ @@ -16,14 +16,14 @@ from jsonschema import Draft202012Validator from policyengine.cli import main as cli_main -from policyengine.core.release_manifest import ( +from policyengine.core.tax_benefit_model import TaxBenefitModel +from policyengine.core.tax_benefit_model_version import TaxBenefitModelVersion +from policyengine.provenance.manifest import ( DataReleaseManifest, get_data_release_manifest, get_release_manifest, ) -from policyengine.core.tax_benefit_model import TaxBenefitModel -from policyengine.core.tax_benefit_model_version import TaxBenefitModelVersion -from policyengine.core.trace_tro import ( +from policyengine.provenance.trace import ( POLICYENGINE_ORGANIZATION, TRACE_TROV_NAMESPACE, build_trace_tro_from_release_bundle, @@ -472,7 +472,7 @@ def test__given_trace_tro_property__then_emits_valid_tro(self): return_value=data_release_manifest, ): with patch( - "policyengine.core.trace_tro.fetch_pypi_wheel_metadata", + "policyengine.provenance.trace.fetch_pypi_wheel_metadata", side_effect=_fake_fetch_pypi, ): tro = model_version.trace_tro @@ -641,7 +641,7 @@ def test__given_trace_tro_stdout__then_writes_canonical_json( return_value=data_release_manifest, ): with patch( - "policyengine.core.trace_tro.fetch_pypi_wheel_metadata", + "policyengine.provenance.trace.fetch_pypi_wheel_metadata", side_effect=_fake_fetch_pypi, ): exit_code = cli_main(["trace-tro", "us"]) @@ -661,7 +661,7 @@ def test__given_out_path__then_writes_to_file(self, tmp_path, monkeypatch): return_value=data_release_manifest, ): with patch( - "policyengine.core.trace_tro.fetch_pypi_wheel_metadata", + "policyengine.provenance.trace.fetch_pypi_wheel_metadata", side_effect=_fake_fetch_pypi, ): exit_code = cli_main(["trace-tro", "us", "--out", str(out)])