Source code for message_ix_models.util.sdmx

"""Utilities for handling objects from :mod:`sdmx`."""
import logging
from datetime import datetime
from enum import Enum, Flag
from importlib.metadata import version
from pathlib import Path
from typing import TYPE_CHECKING, Dict, List, Mapping, Optional, Union

import sdmx
import sdmx.message
from iam_units import registry
from sdmx.model.v21 import AnnotableArtefact, Annotation, Code, InternationalString

from .common import package_data_path

if TYPE_CHECKING:
    from os import PathLike

    import sdmx.model.common

log = logging.getLogger(__name__)

CodeLike = Union[str, Code]


[docs]def as_codes(data: Union[List[str], Dict[str, CodeLike]]) -> List[Code]: """Convert `data` to a :class:`list` of |Code| objects. Various inputs are accepted: - :class:`list` of :class:`str`. - :class:`dict`, in which keys are :attr:`~sdmx.model.common.Code.id` and values are further :class:`dict` with keys matching other Code attributes. """ # Assemble results as a dictionary result: Dict[str, Code] = {} if isinstance(data, list): # FIXME typing ignored temporarily for PR#9 data = dict(zip(data, data)) # type: ignore [arg-type] elif not isinstance(data, Mapping): raise TypeError(data) for id, info in data.items(): # Pass through Code; convert other types to dict() if isinstance(info, Code): result[info.id] = info continue elif isinstance(info, str): _info = dict(name=info) elif isinstance(info, Mapping): _info = dict(info) else: raise TypeError(info) # Create a Code object code = Code( id=str(id), name=_info.pop("name", str(id).title()), ) # Store the description, if any try: code.description = InternationalString(value=_info.pop("description")) except KeyError: pass # Associate with a parent try: parent_id = _info.pop("parent") except KeyError: pass # No parent else: result[parent_id].append_child(code) # Associate with any children for id in _info.pop("child", []): try: code.append_child(result[id]) except KeyError: pass # Not parsed yet # Convert other dictionary (key, value) pairs to annotations for id, value in _info.items(): code.annotations.append( Annotation(id=id, text=value if isinstance(value, str) else repr(value)) ) result[code.id] = code return list(result.values())
[docs]def eval_anno(obj: AnnotableArtefact, id: str): """Retrieve the annotation `id` from `obj`, run :func:`eval` on its contents. .. deprecated:: 2023.9.12 Use :meth:`sdmx.model.common.AnnotableArtefact.eval_annotation`, which provides the same functionality. """ try: value = str(obj.get_annotation(id=id).text) except KeyError: # No such attribute return None try: return eval(value, {"registry": registry}) except Exception as e: # Something that can't be eval()'d, e.g. a plain string log.debug(f"Could not eval({value!r}): {e}") return value
[docs]def make_enum(urn, base=Enum): """Create an :class:`.enum.Enum` (or `base`) with members from codelist `urn`.""" # Read the code list cl = read(urn) # Ensure the 0 member is NONE, not any of the codes names = ["NONE"] if issubclass(base, Flag) else [] names.extend(code.id for code in cl) # Create the class return base(urn, names)
[docs]def read(urn: str, base_dir: Optional["PathLike"] = None): """Read SDMX object from package data given its `urn`.""" # Identify a path that matches `urn` base_dir = Path(base_dir or package_data_path("sdmx")) urn = urn.replace(":", "_") # ":" invalid on Windows paths = sorted( set(base_dir.glob(f"*{urn}*.xml")) | set(base_dir.glob(f"*{urn.upper()}*.xml")) ) if len(paths) > 1: log.info( f"Match {paths[0].relative_to(base_dir)} for {urn!r}; {len(paths) -1 } " "other result(s)" ) try: with open(paths[0], "rb") as f: msg = sdmx.read_sdmx(f) except IndexError: raise FileNotFoundError(f"'*{urn}*.xml', '*{urn.upper()}*.xml' or similar") for _, cls in msg.iter_collections(): try: return next(iter(msg.objects(cls).values())) except StopIteration: pass
[docs]def write(obj, base_dir: Optional["PathLike"] = None, basename: Optional[str] = None): """Store an SDMX object as package data.""" base_dir = Path(base_dir or package_data_path("sdmx")) if isinstance(obj, sdmx.message.StructureMessage): msg = obj assert basename else: # Set the URN of the object obj.urn = sdmx.urn.make(obj) # Wrap the object in a StructureMessage msg = sdmx.message.StructureMessage() msg.add(obj) # Identify a path to write the file. ":" is invalid on Windows. basename = basename or obj.urn.split("=")[-1].replace(":", "_") msg.header = sdmx.message.Header( source=f"Generated by message_ix_models {version('message_ix_models')}", prepared=datetime.now(), ) path = base_dir.joinpath(f"{basename}.xml") # Write path.parent.mkdir(parents=True, exist_ok=True) path.write_bytes(sdmx.to_xml(msg, pretty_print=True)) log.info(f"Wrote {path}")
[docs]def register_agency( agency: "sdmx.model.common.Agency", ) -> "sdmx.model.common.AgencyScheme": """Add `agency` to the :class:`.AgencyScheme` "IIASA_ECE:AGENCIES".""" # Read the existing agency scheme as_ = read("IIASA_ECE:AGENCIES") if agency in as_: log.info(f"Replace or update existing {as_[agency.id]!r}") as_.items[agency.id] = agency else: as_.append(agency) log.info(f"Updated {as_!r}") # Write to file again write(as_) return as_