ModelTrainEval/utils/config.py

"""YAML 配置加载：支持 `_base_` 继承、深度合并与命令行点号覆盖。"""

from __future__ import annotations

import ast
from collections.abc import Iterable, Mapping
from pathlib import Path
from typing import Any

import yaml

BASE_KEY = "_base_"
REPLACE_KEY = "_replace_"


class Config(dict):
    """行为与 dict 一致，额外支持 `cfg.train.lr` 式属性访问。"""

    def __init__(self, *args: Any, **kwargs: Any) -> None:
        super().__init__(*args, **kwargs)
        for key, value in list(self.items()):
            if isinstance(value, Mapping) and not isinstance(value, Config):
                self[key] = Config(value)

    def __getattr__(self, item: str) -> Any:
        try:
            return self[item]
        except KeyError as exc:  # pragma: no cover - 仅用于更友好的报错
            raise AttributeError(item) from exc

    def get_path(self, dotted: str, default: Any = None) -> Any:
        node: Any = self
        for part in dotted.split("."):
            if not isinstance(node, Mapping) or part not in node:
                return default
            node = node[part]
        return node

    def to_dict(self) -> dict:
        def unwrap(value: Any) -> Any:
            if isinstance(value, Mapping):
                return {k: unwrap(v) for k, v in value.items()}
            if isinstance(value, list):
                return [unwrap(v) for v in value]
            return value

        return unwrap(self)

    def dump(self, path: str | Path) -> None:
        Path(path).write_text(
            yaml.safe_dump(self.to_dict(), allow_unicode=True, sort_keys=False),
            encoding="utf-8",
        )


def deep_merge(base: Mapping, override: Mapping) -> dict:
    """递归合并；override 中的标量/列表直接覆盖 base。

    子字典默认是**合并**而不是替换，这在换组件时会出问题：把 backbone 从
    `tiny_conv` 换成 `torchvision`，base 里 tiny_conv 专有的 widths/depths
    会残留下来，传给新组件直接报错。这时在该字典里写 `_replace_: true`，
    表示整块替换掉 base 的同名字典。
    """
    merged = dict(base)
    for key, value in override.items():
        if isinstance(value, Mapping) and value.get(REPLACE_KEY):
            merged[key] = dict(value)  # 整块替换；标记本身留到最后统一剥掉
        elif key in merged and isinstance(merged[key], Mapping) and isinstance(value, Mapping):
            merged[key] = deep_merge(merged[key], value)
        else:
            merged[key] = value
    return merged


def strip_replace_keys(value: Any) -> Any:
    """剥掉所有 `_replace_` 标记，避免它漏进组件的构造参数。

    必须等**全部合并做完**再剥：`_load_raw` 对每个 base 都会先 `deep_merge({}, data)`
    一次，若在那时就剥掉，标记还没来得及参与真正的替换判断就没了。
    """
    if isinstance(value, Mapping):
        return {k: strip_replace_keys(v) for k, v in value.items() if k != REPLACE_KEY}
    if isinstance(value, list):
        return [strip_replace_keys(v) for v in value]
    return value


def _load_raw(path: Path, seen: set[Path]) -> dict:
    path = path.resolve()
    if path in seen:
        raise ValueError(f"配置 `{BASE_KEY}` 出现循环引用: {path}")
    seen.add(path)

    data = yaml.safe_load(path.read_text(encoding="utf-8")) or {}
    if not isinstance(data, dict):
        raise TypeError(f"配置文件必须是映射: {path}")

    bases = data.pop(BASE_KEY, [])
    if isinstance(bases, str):
        bases = [bases]

    merged: dict = {}
    for base in bases:
        base_path = (path.parent / base).resolve()
        merged = deep_merge(merged, _load_raw(base_path, set(seen)))
    return deep_merge(merged, data)


def parse_override(item: str) -> tuple[str, Any]:
    """解析 `train.lr=1e-3`，值按 Python 字面量尝试转换，失败则保留字符串。"""
    if "=" not in item:
        raise ValueError(f"覆盖项格式应为 key=value，收到: {item}")
    key, raw = item.split("=", 1)
    try:
        value = ast.literal_eval(raw)
    except (ValueError, SyntaxError):
        value = raw
    return key.strip(), value


def apply_overrides(cfg: dict, overrides: Iterable[str]) -> dict:
    for item in overrides:
        key, value = parse_override(item)
        parts = key.split(".")
        node = cfg
        for part in parts[:-1]:
            nxt = node.get(part)
            if not isinstance(nxt, dict):
                nxt = {}
                node[part] = nxt
            node = nxt
        node[parts[-1]] = value
    return cfg


def load_config(path: str | Path, overrides: Iterable[str] = ()) -> Config:
    cfg = _load_raw(Path(path), set())
    cfg = apply_overrides(cfg, overrides)
    cfg = strip_replace_keys(cfg)
    cfg.setdefault("config_path", str(Path(path).resolve()))
    return Config(cfg)