bittorrent-benchmarks/benchmarks/core/config.py

"""Basic utilities for structuring experiment configurations based on Pydantic schemas."""
import os
from abc import abstractmethod
from io import TextIOBase
from typing import Type, Dict, TextIO, Callable, cast

import yaml
from typing_extensions import Generic, overload

from benchmarks.core.experiments.experiments import TExperiment
from benchmarks.core.pydantic import SnakeCaseModel


class ExperimentBuilder(SnakeCaseModel, Generic[TExperiment]):
    """:class:`ExperimentBuilders` can build real :class:`Experiment`s out of :class:`ConfigModel`s. """

    @abstractmethod
    def build(self) -> TExperiment:
        pass


class ConfigParser:
    """
    :class:`ConfigParser` is a utility class to parse configuration files into :class:`ExperimentBuilder`s.
    Currently, each :class:`ExperimentBuilder` can appear at most once in the config file.
    """

    def __init__(self):
        self.experiment_types = {}

    def register(self, root: Type[ExperimentBuilder[TExperiment]]):
        self.experiment_types[root.alias()] = root

    @overload
    def parse(self, data: dict) -> Dict[str, ExperimentBuilder[TExperiment]]:
        ...

    @overload
    def parse(self, data: TextIO) -> Dict[str, ExperimentBuilder[TExperiment]]:
        ...

    def parse(self, data: dict | TextIO) -> Dict[str, ExperimentBuilder[TExperiment]]:
        if isinstance(data, TextIOBase):
            entries = yaml.safe_load(os.path.expandvars(data.read()))
        else:
            entries = data

        return {
            tag: self.experiment_types[tag].model_validate(config)
            for tag, config in entries.items()
        }