config_editor.py

import json
import os.path
import sys

from typing import Any
from enum import Enum
from ruamel.yaml import YAML
from loguru import logger


class ConfigType(Enum):
    index = "index"
    tiny = "tiny"
    diffusion = "diffusion"


def edit_config(
        config_type: ConfigType,
        path: str,
        value: Any,
        is_template: bool = False,
        is_num: bool = False,
        is_bool: bool = False
):
    folder = "configs_template" if is_template else "configs"
    name = {
        ConfigType.index: "config",
        ConfigType.tiny: "config_tiny",
        ConfigType.diffusion: "diffusion"
    }[config_type]
    extension = ".json" if config_type is not ConfigType.diffusion else ".yaml"
    config_full_path = folder + "/" + name + ("_template" if is_template else "") + extension

    logger.info(f"Config: {config_full_path}")

    if not os.path.exists(config_full_path):
        raise FileNotFoundError(config_full_path)

    with open(config_full_path, "r+") as f:
        if extension == ".json":
            config = json.load(f)

        elif extension == ".yaml":
            yaml = YAML()
            yaml.preserve_quotes = True
            yaml.indent(offset=2)
            config = yaml.load(f)

        else:
            raise RuntimeError(f"Unexpected extension: {extension}")

        keys = path.split(".")
        current_dict = config

        for key in keys[:-1]:
            if key not in current_dict:
                raise NameError(f"{key} is not in config.")
            current_dict = current_dict[key]

        if is_num:
            if "." in value:
                value = float(value)
                logger.info(f"{value} is a float")
            else:
                value = int(value)
                logger.info(f"{value} is a int")

        if is_bool:
            value = bool(value)

        current_dict[keys[-1]] = value

        f.seek(0)

        if extension == ".json":
            json.dump(config, f, indent=2)
        elif extension == ".yaml":
            yaml.dump(config, f)

        logger.info(f"Done: {path} -> {value}")


def get_value_from_args(key: str, default: Any = None) -> Any:
    if key in sys.argv:
        index = sys.argv.index(key) + 1
        if index < len(sys.argv):
            return sys.argv[index]
    return default


if __name__ == '__main__':
    edit_config(
        ConfigType(
            get_value_from_args("-t", "index")
        ),
        get_value_from_args("-p"),
        get_value_from_args("-v"),
        "--template" in sys.argv or "-T" in sys.argv,
        "--num" in sys.argv or "-N" in sys.argv,
        "--bool" in sys.argv or "-B" in sys.argv
    )