fix: Redo package version management

Author: anth-volk

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"

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
+]

policyengine/utils/versioning.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

tests/fixtures/utils/versioning.py tests/fixtures/utils/packages.pytests/fixtures/utils/versioning.py renamed to tests/fixtures/utils/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)