From ef591ff2721509ceb2b3de28d7319508cda0cae9 Mon Sep 17 00:00:00 2001 From: Anthony Volk Date: Mon, 12 May 2025 20:46:30 -0400 Subject: [PATCH 1/6] feat: Add API version to economy comparisons --- policyengine/constants.py | 23 +++++++ .../calculate_economy_comparison.py | 3 + policyengine/simulation.py | 2 +- policyengine/utils/versioning.py | 30 +++++++++ tests/__init__.py | 0 tests/fixtures/utils/__init__.py | 0 tests/fixtures/utils/versioning.py | 18 ++++++ tests/utils/__init__.py | 0 tests/utils/test_versioning.py | 63 +++++++++++++++++++ 9 files changed, 138 insertions(+), 1 deletion(-) create mode 100644 policyengine/utils/versioning.py create mode 100644 tests/__init__.py create mode 100644 tests/fixtures/utils/__init__.py create mode 100644 tests/fixtures/utils/versioning.py create mode 100644 tests/utils/__init__.py create mode 100644 tests/utils/test_versioning.py diff --git a/policyengine/constants.py b/policyengine/constants.py index c6bca554..68ada5b8 100644 --- a/policyengine/constants.py +++ b/policyengine/constants.py @@ -3,6 +3,29 @@ from policyengine_core.data import Dataset from policyengine.utils.data_download import download +SUPPORTED_COUNTRY_IDS = [ + "us", + "uk", +] + +UNSUPPORTED_COUNTRY_IDS = [ + "ca", + "il", + "ng", +] + +SUPPORTED_COUNTRY_PACKAGES = [ + f"policyengine_{country}" for country in SUPPORTED_COUNTRY_IDS +] +UNSUPPORTED_COUNTRY_PACKAGES = [ + f"policyengine_{country}" if country != "ca" else "policyengine_canada" + for country in UNSUPPORTED_COUNTRY_IDS +] + +ALL_COUNTRY_PACKAGES = ( + SUPPORTED_COUNTRY_PACKAGES + UNSUPPORTED_COUNTRY_PACKAGES +) + # Datasets ENHANCED_FRS = "hf://policyengine/policyengine-uk-data/enhanced_frs_2022_23.h5" diff --git a/policyengine/outputs/macro/comparison/calculate_economy_comparison.py b/policyengine/outputs/macro/comparison/calculate_economy_comparison.py index 6592c305..8b6ee3d0 100644 --- a/policyengine/outputs/macro/comparison/calculate_economy_comparison.py +++ b/policyengine/outputs/macro/comparison/calculate_economy_comparison.py @@ -10,6 +10,7 @@ from policyengine.outputs.macro.single.calculate_single_economy import ( SingleEconomy, ) +from policyengine.utils.versioning import get_country_package_version from typing import List, Dict @@ -775,6 +776,7 @@ def uk_constituency_breakdown( class EconomyComparison(BaseModel): + country_packge_version: str budget: BudgetaryImpact detailed_budget: DetailedBudgetaryImpact decile: DecileImpact @@ -823,6 +825,7 @@ def calculate_economy_comparison( ) return EconomyComparison( + country_package_version=get_country_package_version(country_id), budget=budgetary_impact_data, detailed_budget=detailed_budgetary_impact_data, decile=decile_impact_data, diff --git a/policyengine/simulation.py b/policyengine/simulation.py index 61f00acb..ecb794fc 100644 --- a/policyengine/simulation.py +++ b/policyengine/simulation.py @@ -181,7 +181,7 @@ def _initialise_simulation( time_period: TimePeriodType, region: RegionType, subsample: SubsampleType, - ): + ) -> CountrySimulation: macro = scope == "macro" _simulation_type: Type[CountrySimulation] = { "uk": { diff --git a/policyengine/utils/versioning.py b/policyengine/utils/versioning.py new file mode 100644 index 00000000..0eb27491 --- /dev/null +++ b/policyengine/utils/versioning.py @@ -0,0 +1,30 @@ +from importlib.metadata import version +from policyengine.constants import ( + SUPPORTED_COUNTRY_IDS, + UNSUPPORTED_COUNTRY_IDS, +) + + +def get_version(package: str) -> str: + try: + return version(package) + except Exception as e: + raise RuntimeError( + f"Could not get version for package {package}: {e}" + ) from e + + +def get_country_package_version(country_id: str) -> str: + + if country_id in UNSUPPORTED_COUNTRY_IDS: + raise ValueError(f"Country ID {country_id} is not supported.") + + if country_id not in SUPPORTED_COUNTRY_IDS: + raise ValueError(f"Country ID {country_id} is not recognized.") + + try: + return version(f"policyengine_{country_id}") + except Exception as e: + raise RuntimeError( + f"Could not get version for package policyengine_{country_id}: {e}" + ) from e diff --git a/tests/__init__.py b/tests/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/tests/fixtures/utils/__init__.py b/tests/fixtures/utils/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/tests/fixtures/utils/versioning.py b/tests/fixtures/utils/versioning.py new file mode 100644 index 00000000..fd9d81da --- /dev/null +++ b/tests/fixtures/utils/versioning.py @@ -0,0 +1,18 @@ +from policyengine.constants import SUPPORTED_COUNTRY_PACKAGES +from importlib.metadata import PackageNotFoundError +import pytest +from unittest.mock import patch + + +@pytest.fixture +def patch_importlib_version(): + def mock_version(package_name): + if package_name in SUPPORTED_COUNTRY_PACKAGES: + return "1.0.0" + else: + raise Exception(f"Package {package_name} not found") + + with patch( + "policyengine.utils.versioning.version", side_effect=mock_version + ) as mock: + yield mock diff --git a/tests/utils/__init__.py b/tests/utils/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/tests/utils/test_versioning.py b/tests/utils/test_versioning.py new file mode 100644 index 00000000..f7fc450b --- /dev/null +++ b/tests/utils/test_versioning.py @@ -0,0 +1,63 @@ +import pytest +from policyengine.utils.versioning import ( + get_version, + get_country_package_version, +) +from tests.fixtures.utils.versioning import patch_importlib_version +from importlib.metadata import PackageNotFoundError + + +class TestGetVersion: + def test__given_package_exist__then_return_version( + self, patch_importlib_version + ): + test_package = "policyengine_us" + + test_version = get_version(test_package) + + # Version number defined by mock + assert test_version == "1.0.0" + + def test__given_package_does_not_exist__then_raise_exception( + self, patch_importlib_version + ): + test_package = "non_existent_package" + + with pytest.raises( + RuntimeError, + match=f"Could not get version for package {test_package}", + ): + get_version(test_package) + + +class TestGetCountryPackageVersion: + def test__given_country_package_exists__then_return_version( + self, patch_importlib_version + ): + test_country = "us" + + test_version = get_country_package_version(test_country) + + # Version number defined by mock + assert test_version == "1.0.0" + + def test__given_country_package_does_not_exist__then_raise_exception( + self, patch_importlib_version + ): + test_country = "non_existent_country" + + with pytest.raises( + ValueError, match=f"Country ID {test_country} is not recognized." + ): + get_country_package_version(test_country) + + def test__given_country_not_supported__then_raise_error( + self, patch_importlib_version + ): + unsupported_country = "ca" + + with pytest.raises( + ValueError, + match=f"Country ID {unsupported_country} is not supported.", + ): + get_country_package_version(unsupported_country) From ad6ac1ac9dc6e146a550287385fc020d2368fd34 Mon Sep 17 00:00:00 2001 From: Anthony Volk Date: Tue, 13 May 2025 13:20:05 -0400 Subject: [PATCH 2/6] fix: Fix typo in EconomyComparison schema --- .../outputs/macro/comparison/calculate_economy_comparison.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/policyengine/outputs/macro/comparison/calculate_economy_comparison.py b/policyengine/outputs/macro/comparison/calculate_economy_comparison.py index 8b6ee3d0..176b4d5b 100644 --- a/policyengine/outputs/macro/comparison/calculate_economy_comparison.py +++ b/policyengine/outputs/macro/comparison/calculate_economy_comparison.py @@ -776,7 +776,7 @@ def uk_constituency_breakdown( class EconomyComparison(BaseModel): - country_packge_version: str + country_package_version: str budget: BudgetaryImpact detailed_budget: DetailedBudgetaryImpact decile: DecileImpact From d8e3165b9751ee22f5c0e0a14244c668abbe869c Mon Sep 17 00:00:00 2001 From: Anthony Volk Date: Tue, 13 May 2025 13:22:52 -0400 Subject: [PATCH 3/6] fix: Reformat package name constants --- policyengine/constants.py | 13 +++++++++++-- 1 file changed, 11 insertions(+), 2 deletions(-) diff --git a/policyengine/constants.py b/policyengine/constants.py index 68ada5b8..1ea474bf 100644 --- a/policyengine/constants.py +++ b/policyengine/constants.py @@ -17,9 +17,18 @@ SUPPORTED_COUNTRY_PACKAGES = [ f"policyengine_{country}" for country in SUPPORTED_COUNTRY_IDS ] + + +def _package_name_for(country_id: str) -> str: + return ( + f"policyengine_{country_id}" + if country_id != "ca" + else "policyengine_canada" + ) + + UNSUPPORTED_COUNTRY_PACKAGES = [ - f"policyengine_{country}" if country != "ca" else "policyengine_canada" - for country in UNSUPPORTED_COUNTRY_IDS + _package_name_for(country) for country in UNSUPPORTED_COUNTRY_IDS ] ALL_COUNTRY_PACKAGES = ( From 9f2ee8ba7da7a2c33de1cdd7fc7fd6520e16eba9 Mon Sep 17 00:00:00 2001 From: Anthony Volk Date: Tue, 13 May 2025 13:43:37 -0400 Subject: [PATCH 4/6] fix: Update versioning function and tests --- policyengine/utils/versioning.py | 24 +++----------- tests/fixtures/utils/versioning.py | 8 +++-- tests/utils/test_versioning.py | 51 ++++++------------------------ 3 files changed, 19 insertions(+), 64 deletions(-) diff --git a/policyengine/utils/versioning.py b/policyengine/utils/versioning.py index 0eb27491..b4c490a9 100644 --- a/policyengine/utils/versioning.py +++ b/policyengine/utils/versioning.py @@ -5,26 +5,10 @@ ) -def get_version(package: str) -> str: - try: - return version(package) - except Exception as e: - raise RuntimeError( - f"Could not get version for package {package}: {e}" - ) from e - - def get_country_package_version(country_id: str) -> str: - if country_id in UNSUPPORTED_COUNTRY_IDS: - raise ValueError(f"Country ID {country_id} is not supported.") - - if country_id not in SUPPORTED_COUNTRY_IDS: - raise ValueError(f"Country ID {country_id} is not recognized.") + # Canada package doesn't use two-letter country code + if country_id == "ca": + country_id = "canada" - try: - return version(f"policyengine_{country_id}") - except Exception as e: - raise RuntimeError( - f"Could not get version for package policyengine_{country_id}: {e}" - ) from e + return version(f"policyengine_{country_id}") diff --git a/tests/fixtures/utils/versioning.py b/tests/fixtures/utils/versioning.py index fd9d81da..bf3df885 100644 --- a/tests/fixtures/utils/versioning.py +++ b/tests/fixtures/utils/versioning.py @@ -1,14 +1,16 @@ -from policyengine.constants import SUPPORTED_COUNTRY_PACKAGES +from policyengine.constants import ALL_COUNTRY_PACKAGES from importlib.metadata import PackageNotFoundError import pytest from unittest.mock import patch +MOCK_VERSION = "MOCK_VERSION" + @pytest.fixture def patch_importlib_version(): def mock_version(package_name): - if package_name in SUPPORTED_COUNTRY_PACKAGES: - return "1.0.0" + if package_name in ALL_COUNTRY_PACKAGES: + return MOCK_VERSION else: raise Exception(f"Package {package_name} not found") diff --git a/tests/utils/test_versioning.py b/tests/utils/test_versioning.py index f7fc450b..69454654 100644 --- a/tests/utils/test_versioning.py +++ b/tests/utils/test_versioning.py @@ -1,37 +1,16 @@ import pytest from policyengine.utils.versioning import ( - get_version, get_country_package_version, ) -from tests.fixtures.utils.versioning import patch_importlib_version +from tests.fixtures.utils.versioning import ( + patch_importlib_version, + MOCK_VERSION, +) from importlib.metadata import PackageNotFoundError -class TestGetVersion: - def test__given_package_exist__then_return_version( - self, patch_importlib_version - ): - test_package = "policyengine_us" - - test_version = get_version(test_package) - - # Version number defined by mock - assert test_version == "1.0.0" - - def test__given_package_does_not_exist__then_raise_exception( - self, patch_importlib_version - ): - test_package = "non_existent_package" - - with pytest.raises( - RuntimeError, - match=f"Could not get version for package {test_package}", - ): - get_version(test_package) - - class TestGetCountryPackageVersion: - def test__given_country_package_exists__then_return_version( + def test__given_package_exists__then_return_version( self, patch_importlib_version ): test_country = "us" @@ -39,25 +18,15 @@ def test__given_country_package_exists__then_return_version( test_version = get_country_package_version(test_country) # Version number defined by mock - assert test_version == "1.0.0" + assert test_version == MOCK_VERSION - def test__given_country_package_does_not_exist__then_raise_exception( + def test__given_package_does_not_exist__then_raise_exception( self, patch_importlib_version ): - test_country = "non_existent_country" + test_country = "zz" with pytest.raises( - ValueError, match=f"Country ID {test_country} is not recognized." + Exception, + match=f"Package policyengine_{test_country} not found", ): get_country_package_version(test_country) - - def test__given_country_not_supported__then_raise_error( - self, patch_importlib_version - ): - unsupported_country = "ca" - - with pytest.raises( - ValueError, - match=f"Country ID {unsupported_country} is not supported.", - ): - get_country_package_version(unsupported_country) From 642952b2482b4b01922896376416a2bfcf4d9dc5 Mon Sep 17 00:00:00 2001 From: Anthony Volk Date: Tue, 13 May 2025 15:58:12 -0400 Subject: [PATCH 5/6] fix: Redo package version management --- policyengine/constants.py | 33 ------------ policyengine/utils/packages.py | 28 ++++++++++ policyengine/utils/versioning.py | 14 ----- .../utils/{versioning.py => packages.py} | 7 ++- tests/utils/test_packages.py | 54 +++++++++++++++++++ tests/utils/test_versioning.py | 32 ----------- 6 files changed, 85 insertions(+), 83 deletions(-) create mode 100644 policyengine/utils/packages.py delete mode 100644 policyengine/utils/versioning.py rename tests/fixtures/utils/{versioning.py => packages.py} (58%) create mode 100644 tests/utils/test_packages.py delete mode 100644 tests/utils/test_versioning.py diff --git a/policyengine/constants.py b/policyengine/constants.py index 1ea474bf..43351271 100644 --- a/policyengine/constants.py +++ b/policyengine/constants.py @@ -3,40 +3,7 @@ from policyengine_core.data import Dataset from policyengine.utils.data_download import download -SUPPORTED_COUNTRY_IDS = [ - "us", - "uk", -] - -UNSUPPORTED_COUNTRY_IDS = [ - "ca", - "il", - "ng", -] - -SUPPORTED_COUNTRY_PACKAGES = [ - f"policyengine_{country}" for country in SUPPORTED_COUNTRY_IDS -] - - -def _package_name_for(country_id: str) -> str: - return ( - f"policyengine_{country_id}" - if country_id != "ca" - else "policyengine_canada" - ) - - -UNSUPPORTED_COUNTRY_PACKAGES = [ - _package_name_for(country) for country in UNSUPPORTED_COUNTRY_IDS -] - -ALL_COUNTRY_PACKAGES = ( - SUPPORTED_COUNTRY_PACKAGES + UNSUPPORTED_COUNTRY_PACKAGES -) - # Datasets - ENHANCED_FRS = "hf://policyengine/policyengine-uk-data/enhanced_frs_2022_23.h5" FRS = "hf://policyengine/policyengine-uk-data/frs_2022_23.h5" ENHANCED_CPS = "hf://policyengine/policyengine-us-data/enhanced_cps_2024.h5" diff --git a/policyengine/utils/packages.py b/policyengine/utils/packages.py new file mode 100644 index 00000000..78a254f2 --- /dev/null +++ b/policyengine/utils/packages.py @@ -0,0 +1,28 @@ +from importlib.metadata import version + + +def get_country_package_name(country_id: str) -> str: + if country_id in NON_STANDARD_COUNTRY_CODES: + return f"policyengine_{NON_STANDARD_COUNTRY_CODES[country_id]}" + if country_id in COUNTRY_IDS: + return f"policyengine_{country_id}" + raise ValueError( + f"Unsupported country ID: {country_id}. Supported IDs are: {COUNTRY_IDS}" + ) + + +def get_country_package_version(country_id: str) -> str: + + package_name = get_country_package_name(country_id) + return version(package_name) + + +COUNTRY_IDS = ["us", "uk", "ca", "il", "ng"] + +NON_STANDARD_COUNTRY_CODES = { + "ca": "canada", +} + +COUNTRY_PACKAGES = [ + get_country_package_name(country) for country in COUNTRY_IDS +] diff --git a/policyengine/utils/versioning.py b/policyengine/utils/versioning.py deleted file mode 100644 index b4c490a9..00000000 --- a/policyengine/utils/versioning.py +++ /dev/null @@ -1,14 +0,0 @@ -from importlib.metadata import version -from policyengine.constants import ( - SUPPORTED_COUNTRY_IDS, - UNSUPPORTED_COUNTRY_IDS, -) - - -def get_country_package_version(country_id: str) -> str: - - # Canada package doesn't use two-letter country code - if country_id == "ca": - country_id = "canada" - - return version(f"policyengine_{country_id}") diff --git a/tests/fixtures/utils/versioning.py b/tests/fixtures/utils/packages.py similarity index 58% rename from tests/fixtures/utils/versioning.py rename to tests/fixtures/utils/packages.py index bf3df885..2877aef7 100644 --- a/tests/fixtures/utils/versioning.py +++ b/tests/fixtures/utils/packages.py @@ -1,5 +1,4 @@ -from policyengine.constants import ALL_COUNTRY_PACKAGES -from importlib.metadata import PackageNotFoundError +from policyengine.utils.packages import COUNTRY_PACKAGES import pytest from unittest.mock import patch @@ -9,12 +8,12 @@ @pytest.fixture def patch_importlib_version(): def mock_version(package_name): - if package_name in ALL_COUNTRY_PACKAGES: + if package_name in COUNTRY_PACKAGES: return MOCK_VERSION else: raise Exception(f"Package {package_name} not found") with patch( - "policyengine.utils.versioning.version", side_effect=mock_version + "policyengine.utils.packages.version", side_effect=mock_version ) as mock: yield mock diff --git a/tests/utils/test_packages.py b/tests/utils/test_packages.py new file mode 100644 index 00000000..822377b0 --- /dev/null +++ b/tests/utils/test_packages.py @@ -0,0 +1,54 @@ +import pytest +from policyengine.utils.packages import ( + get_country_package_name, + get_country_package_version, +) +from tests.fixtures.utils.packages import ( + patch_importlib_version, + MOCK_VERSION, +) + + +class TestGetCountryPackageName: + def test__given_country_id__then_return_package_name(self): + test_country = "us" + + test_package_name = get_country_package_name(test_country) + + assert test_package_name == "policyengine_us" + + def test__given_non_standard_country_id__then_return_package_name(self): + test_country = "ca" + + test_package_name = get_country_package_name(test_country) + + assert test_package_name == "policyengine_canada" + + def test__given_unsupported_country_id__then_return_package_name(self): + test_country = "zz" + + with pytest.raises(Exception, match="Unsupported country ID: zz"): + get_country_package_name(test_country) + + +class TestGetCountryPackageVersion: + def test__given_package_exists__then_return_version( + self, patch_importlib_version + ): + test_country = "us" + + test_version = get_country_package_version(test_country) + + # Version number defined by mock + assert test_version == MOCK_VERSION + + def test__given_package_does_not_exist__then_raise_exception( + self, patch_importlib_version + ): + test_country = "zz" + + with pytest.raises( + Exception, + match=f"Unsupported country ID: {test_country}. Supported IDs are: ", + ): + get_country_package_version(test_country) diff --git a/tests/utils/test_versioning.py b/tests/utils/test_versioning.py deleted file mode 100644 index 69454654..00000000 --- a/tests/utils/test_versioning.py +++ /dev/null @@ -1,32 +0,0 @@ -import pytest -from policyengine.utils.versioning import ( - get_country_package_version, -) -from tests.fixtures.utils.versioning import ( - patch_importlib_version, - MOCK_VERSION, -) -from importlib.metadata import PackageNotFoundError - - -class TestGetCountryPackageVersion: - def test__given_package_exists__then_return_version( - self, patch_importlib_version - ): - test_country = "us" - - test_version = get_country_package_version(test_country) - - # Version number defined by mock - assert test_version == MOCK_VERSION - - def test__given_package_does_not_exist__then_raise_exception( - self, patch_importlib_version - ): - test_country = "zz" - - with pytest.raises( - Exception, - match=f"Package policyengine_{test_country} not found", - ): - get_country_package_version(test_country) From 17b7ad8bd7e67dba02266c3581055d57d8802225 Mon Sep 17 00:00:00 2001 From: Anthony Volk Date: Tue, 13 May 2025 16:02:17 -0400 Subject: [PATCH 6/6] fix: Fix import --- .../outputs/macro/comparison/calculate_economy_comparison.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/policyengine/outputs/macro/comparison/calculate_economy_comparison.py b/policyengine/outputs/macro/comparison/calculate_economy_comparison.py index 176b4d5b..f53691f9 100644 --- a/policyengine/outputs/macro/comparison/calculate_economy_comparison.py +++ b/policyengine/outputs/macro/comparison/calculate_economy_comparison.py @@ -10,7 +10,7 @@ from policyengine.outputs.macro.single.calculate_single_economy import ( SingleEconomy, ) -from policyengine.utils.versioning import get_country_package_version +from policyengine.utils.packages import get_country_package_version from typing import List, Dict