Source code for genno.compat.pyam

import logging
from functools import partial
from importlib import import_module
from importlib.util import find_spec
from typing import TYPE_CHECKING

from genno import Computer, Key
from genno.core.key import single_key

if TYPE_CHECKING:
    from collections.abc import MutableMapping

log = logging.getLogger(__name__)

#: :class:`bool` indicating whether :mod:`pyam` is available. If either the package is
#: not installed, or it is installed but raises some exception on import, this will be
#: :any:`False`, and the operators and configuration handling in the current module will
#: be unavailable.
HAS_PYAM = False

if find_spec("pyam"):
    try:
        import_module("pyam")
    except Exception as e:
        # Handle pyam is installed but cannot be imported
        log.warning(f"{__name__} unavailable due to {e}")
    else:
        HAS_PYAM = True




def handle_config(c: Computer, info: "MutableMapping") -> None:
    """Handle one entry from the ``iamc:`` config section."""
    try:
        c.require_compat("pyam")
    except ModuleNotFoundError:  # pragma: no cover
        if not HAS_PYAM:
            log.warning("Missing pyam; configuration section 'iamc:' ignored")
            return
        else:
            raise

    from . import util

    # For each quantity, use a chain of computations to prepare it
    name = info.pop("variable")

    # Chain of keys produced: first entry is the key for the base quantity
    keys: list[Key] = [Key(info.pop("base"))]

    # Second entry is a simple rename
    keys.append(single_key(c.add_single(Key(name, keys[0].dims, keys[0].tag), keys[0])))

    # Optionally select a subset of data from the base quantity
    sel = info.get("select")
    if sel:
        keys.append(
            single_key(
                c.add_single(
                    keys[-1].add_tag("sel"),
                    (c.get_operator("select"), keys[-1], sel),
                    strict=True,
                )
            )
        )

    # Use a setting for the collapse callback function. This doesn't work from file,
    # since no way to define a Python function in JSON or YAML
    collapse_info = info.pop("collapse", {})
    collapse_func = collapse_info.pop("callback", util.collapse)

    # Use the Computer method to add the conversion step
    # NB convert_pyam() returns a single key when applied to a single key
    keys.append(
        single_key(
            c.add(
                keys[-1],
                "as_pyam",
                rename=info.pop("rename", {}),
                collapse=partial(collapse_func, **collapse_info),
                replace=info.pop("replace", {}),
                drop=set(info.pop("drop", [])) & set(keys[-1].dims),
                unit=info.pop("unit", None),
            )
        )
    )

    log.info(f"Add {repr(keys[-1])} from {repr(keys[0])}")
    log.debug(f"    {len(keys)} keys total")



if HAS_PYAM:
    # Register the configuration handler only if pyam is actually available
    import genno.config

    genno.config.handles("iamc")(handle_config)