Source code for idmtools.builders.yaml_simulation_builder

"""
idmtools YamlSimulationBuilder definition.

Copyright 2021, Bill & Melinda Gates Foundation. All rights reserved.
"""
from logging import getLogger
from typing import Union, Callable, Dict, Any
import yaml
from idmtools.builders import ArmSimulationBuilder
from idmtools.builders.arm_simulation_builder import SweepArm, ArmType

logger = getLogger(__name__)


class DefaultParamFuncDict(dict):
    """
    Enables a function that takes a single parameter and return another function.

    Notes:
        TODO Add Example and types
    """

    def __init__(self, default):
        """
        Initialize our DefaultParamFuncDict.

        Args:
            default: Default function to use
        """
        super().__init__()
        self.__default = default

    def __getitem__(self, item):
        """
        Get item from the DefaultParamFuncDict. It proxies most calls to the function we wrap.

        Args:
            item: Item to lookup

        Returns:
            None
        """
        if item in self:
            return super().__getitem__(item)
        else:
            return self.__default(item)


class YamlSimulationBuilder(ArmSimulationBuilder):
    """
    Class that represents an experiment builder.

    Examples:
        .. literalinclude:: ../../examples/builders/yaml_builder_python.py
    """
    def __init__(self):
        """Constructor."""
        super().__init__()

    def add_sweeps_from_file(self, file_path, func_map: Union[Dict[str, Callable], Callable[[Any], Dict]] = None, sweep_type=ArmType.cross):
        """
        Add sweeps from a file.

        Args:
            file_path: Path to file
            func_map: Optional function map
            sweep_type: Type of sweep

        Returns:
            None
        """
        if func_map is None:
            func_map = {}
        # if the user passing a single function, map it to all values
        elif callable(func_map):
            func_map = DefaultParamFuncDict(func_map)
        with open(file_path, 'r') as stream:
            try:
                parsed = yaml.safe_load(stream)
            except yaml.YAMLError as exc:
                logger.exception(exc)
                exit()

        d_funcs = parsed.values()

        for sweeps in list(d_funcs):
            sweeps_tuples = ((list(d.keys())[0], list(d.values())[0]) for d in sweeps)
            funcs = []
            for func, values in sweeps_tuples:
                funcs.append((func_map[func], values))

            arm = SweepArm(sweep_type, funcs)
            self.add_arm(arm)

    def __iter__(self):
        """
        Iterator the sweep definitions.

        Returns:
            Iterator
        """
        yield from self.sweep_definitions