feat: Add API version to economy comparisons

Author: anth-volk

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"

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,

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": {

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

tests/__init__.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

tests/fixtures/utils/__init__.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)