src/easytranslate/utils/config.py

"""
Configuration helpers (shared module).

Load YAML, merge overrides, apply CLI dotlist overrides, save to disk.
"""

from __future__ import annotations

import json
import logging
from pathlib import Path
from typing import Any

from omegaconf import DictConfig, OmegaConf, open_dict

logger = logging.getLogger(__name__)


def load_config(config_path: str | Path) -> DictConfig:
    """
    Load a YAML file into an OmegaConf DictConfig.

    Args:
        config_path: Path to the YAML file.

    Returns:
        DictConfig instance.
    """
    config_path = Path(config_path)
    if not config_path.exists():
        raise FileNotFoundError(f"Config file not found: {config_path}")

    config = OmegaConf.load(config_path)
    logger.info("Loaded config from %s", config_path)
    return config


def merge_configs(base_config: DictConfig, override_config: DictConfig) -> DictConfig:
    """Merge two configs; values in override win."""
    return OmegaConf.merge(base_config, override_config)


def config_from_cli(config_path: str, cli_args: list[str]) -> DictConfig:
    """
    Load YAML then apply OmegaConf CLI overrides.

    Example:
        python train.py --config configs/default.yaml training.lr=1e-4 model.nhead=16

    Args:
        config_path: Base YAML path.
        cli_args: Dotlist overrides, e.g. ``["training.lr=1e-4"]``.

    Returns:
        Merged DictConfig.
    """
    config = load_config(config_path)

    if cli_args:
        cli_config = OmegaConf.from_cli(cli_args)
        config = merge_configs(config, cli_config)
        logger.info("Applied %d CLI overrides", len(cli_args))

    return config


def overrides_to_cli_args(overrides: dict[str, Any]) -> list[str]:
    """
    Turn a flat dict of dotted keys into OmegaConf CLI strings.

    Example:
        {"model.type": "finetune_nllb", "training.epochs": 2}
        -> ["model.type=finetune_nllb", "training.epochs=2"]
    """
    cli: list[str] = []
    for key, value in overrides.items():
        if not isinstance(key, str):
            raise TypeError(f"Override keys must be str, got {type(key)}")
        cli.append(f"{key}={_format_cli_value(value)}")
    return cli


def _format_cli_value(value: Any) -> str:
    if isinstance(value, bool):
        return str(value).lower()
    if value is None:
        return "null"
    if isinstance(value, (list, dict)):
        return json.dumps(value, ensure_ascii=False)
    return str(value)


def apply_nested_overrides(config: DictConfig, overrides: dict[str, Any]) -> DictConfig:
    """Apply dotted-path updates in-place (missing intermediate keys are created)."""
    with open_dict(config):
        for path, value in overrides.items():
            OmegaConf.update(config, str(path), value, merge=True)
    return config


def save_config(config: DictConfig, path: str | Path) -> None:
    """Save DictConfig to a YAML file."""
    path = Path(path)
    path.parent.mkdir(parents=True, exist_ok=True)
    OmegaConf.save(config, path)


def config_to_dict(config: DictConfig) -> dict:
    """Convert DictConfig to a plain Python dict (with interpolation resolved)."""
    return OmegaConf.to_container(config, resolve=True)