Convert ISO 3166-1 to SDMX

Author: khaeru

transport_data/iso/__init__.py

@@ -1,11 +1,17 @@
 """International Organization for Standardization (ISO)."""
 
 import logging
+from functools import lru_cache
+from pathlib import Path
+from typing import TYPE_CHECKING, Literal, Mapping
 
 from sdmx.model import common
 
 from transport_data.util.pluggy import hookimpl
 
+if TYPE_CHECKING:
+    import gettext
+
 log = logging.getLogger(__name__)
 
 
@@ -18,3 +24,87 @@ def get_agencies():
         contact=[common.Contact(uri=["https://iso.org"])],
     )
     return (a,)
+
+
+def get_cl_iso_3166_1(
+    id_field: Literal["alpha_2", "alpha_3", "numeric"] = "alpha_2",
+) -> common.Codelist:
+    """Generate a :class:`~sdmx.model.common.Codelist` with entries from ISO 3166-1.
+
+    Codes have:
+
+    - IDs according to `id_field`.
+    - An :attr:`~sdmx.model.common.IdentifiableArtefact.name` attribute localized to all
+      the languages present in the upstream database.
+
+    Parameter
+    ---------
+    id_field :
+        Field from the database to use for the IDs of generated Codes.
+    """
+    from importlib.metadata import version
+
+    import sdmx.urn
+    from pycountry import countries as db
+
+    from transport_data import STORE
+
+    # Create an empty codelist
+    cl: common.Codelist = common.Codelist(
+        id=f"{db.root_key}_{id_field}",
+        maintainer=get_agencies()[0],
+        version=version("pycountry"),
+    )
+
+    # Load localizations of this code list
+    translations = load_translations(f"iso{db.root_key}")
+
+    # Convert all entries in the database to SDMX Codes
+    for data in db:
+        # Collect localizations of the country name
+        name = localize_all(data.name, translations)
+        # TODO Collect localizations of other fields; add as annotations
+        # TODO Annotate with other non-localized fields
+
+        # Create a Code
+        c = common.Code(id=getattr(data, id_field), name=name)
+        # Append to the code list
+        cl.append(c)
+
+        # Generate its URN
+        c.urn = sdmx.urn.make(c)
+
+    # Write to local store
+    STORE.set(cl)
+
+    return cl
+
+
+def localize_all(value: str, translations, *, default_locale="en") -> dict[str, str]:
+    """Localize `value` in all languages available in `translations`."""
+    # Put the default locale first
+    result = {default_locale: value}
+
+    for lang, tr in translations.items():
+        localized = tr.gettext(value)
+        if localized != value:
+            result[lang] = localized
+
+    return result
+
+
+@lru_cache
+def load_translations(domain: str) -> Mapping[str, "gettext.NullTranslations"]:
+    """Load all available :mod:`pycountry` translations for `domain`."""
+    from gettext import translation
+
+    from pycountry import LOCALES_DIR
+
+    result = {}
+    for lang in map(lambda d: d.name, Path(LOCALES_DIR).iterdir()):
+        try:
+            result[lang] = translation(domain, LOCALES_DIR, languages=[lang])
+        except FileNotFoundError:
+            pass
+
+    return result

transport_data/tests/test_iso.py

@@ -0,0 +1,21 @@
+import pytest
+
+from transport_data.iso import get_cl_iso_3166_1
+
+
+@pytest.mark.parametrize(
+    "id_field, entries",
+    (
+        ("alpha_2", ("AW", "ZE")),
+        ("alpha_3", ("ABW", "ZWE")),
+        ("numeric", ("533", "716")),
+    ),
+)
+def test_get_cl_iso_3166_1(id_field, entries) -> None:
+    result = get_cl_iso_3166_1(id_field=id_field)
+
+    # Result has the expected number of codes
+    assert 249 == len(result)
+
+    # The name of Aruba is localized in 58 languages
+    assert 58 == len(result[entries[0]].name.localizations)