chore: type version registry metadata

Author: anth-volk

projects/policyengine-api-simulation/src/modal/utils/update_version_registry.py

@@ -24,13 +24,33 @@
 import argparse
 import modal
 from packaging.version import InvalidVersion, Version
+from typing import TypedDict
 
 POLICYENGINE_VERSION_DICT_NAME = "simulation-api-policyengine-versions"
 US_VERSION_DICT_NAME = "simulation-api-us-versions"
 UK_VERSION_DICT_NAME = "simulation-api-uk-versions"
 APP_RELEASE_BUNDLES_DICT_NAME = "simulation-api-app-release-bundles"
 
 
+class CountryBundleMetadata(TypedDict):
+    country: str
+    model_package_name: str
+    model_version: str
+    data_package_name: str
+    data_version: str
+    default_dataset: str
+    default_dataset_uri: str
+    dataset_uris: dict[str, str]
+    dataset_aliases: dict[str, str]
+
+
+class AppReleaseBundleMetadata(TypedDict):
+    app_name: str
+    policyengine_version: str
+    us: CountryBundleMetadata
+    uk: CountryBundleMetadata
+
+
 def _is_newer_version(candidate: str, current: str | None) -> bool:
     """Return True when ``candidate`` should replace ``current`` as 'latest'.
 
@@ -116,7 +136,7 @@ def update_version_dict(
         )
 
 
-def _country_bundle_metadata(country: str) -> dict:
+def _country_bundle_metadata(country: str) -> CountryBundleMetadata:
     from policyengine_api_simulation.release_bundle import (
         DATASET_ALIASES,
         get_country_release_bundle,
@@ -140,7 +160,7 @@ def build_app_release_bundle_metadata(
     *,
     app_name: str,
     policyengine_version: str,
-) -> dict:
+) -> AppReleaseBundleMetadata:
     return {
         "app_name": app_name,
         "policyengine_version": policyengine_version,

projects/policyengine-api-simulation/tests/test_update_version_registry.py

@@ -140,11 +140,23 @@ def test_put_app_release_bundle_metadata_records_app_and_py_version_aliases(
     patched_modal,
     monkeypatch,
 ):
-    def fake_country_bundle_metadata(country: str) -> dict:
+    def fake_country_bundle_metadata(
+        country: str,
+    ) -> registry.CountryBundleMetadata:
         return {
             "country": country,
+            "model_package_name": (
+                "policyengine-us" if country == "us" else "policyengine-uk"
+            ),
             "model_version": "1.0.0" if country == "us" else "2.0.0",
+            "data_package_name": (
+                "policyengine-us-data"
+                if country == "us"
+                else "policyengine-uk-data"
+            ),
             "data_version": "3.0.0" if country == "us" else "4.0.0",
+            "default_dataset": "default",
+            "default_dataset_uri": f"hf://datasets/policyengine/{country}/default",
             "dataset_uris": {"default": f"hf://datasets/policyengine/{country}"},
             "dataset_aliases": {"alias": "default"},
         }