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 字面量解析（`1e-3`、`[5, 6]`、`None`、`False` 都能认），
    再退回 YAML 标量解析，好让配置文件里的写法在命令行上同样可用
    （`null`、`true`、`false`）。两者都失败才保留字符串。

    这一步不能只用 `ast.literal_eval`：YAML 风格的 `false` 会原样变成字符串，
    而非空字符串是**真值**——`--set xxx.enabled=false` 反而把开关打开了。
    也不能只用 YAML：PyYAML 把 `1e-3` 当字符串（YAML 1.1 要写成 `1.0e-3`）。
    """
    if "=" not in item:
        raise ValueError(f"覆盖项格式应为 key=value，收到: {item}")
    key, raw = item.split("=", 1)
    raw = raw.strip()

    try:
        return key.strip(), ast.literal_eval(raw)
    except (ValueError, SyntaxError):
        pass
    try:
        parsed = yaml.safe_load(raw)
    except yaml.YAMLError:
        return key.strip(), raw
    # YAML 会把 "abc: def" 之类解析成 dict，命令行上应当按字符串处理
    return key.strip(), parsed if isinstance(parsed, (bool, int, float, list, type(None))) else raw


def apply_overrides(cfg: dict, overrides: Iterable[str], strict: bool = True) -> dict:
    """把 `a.b.c=value` 写进配置。

    `strict` 时只允许覆盖**已存在**的键：`--set` 打错一个字母（optimizer -> optmizer）
    会凭空造出一个谁也不读的键，训练照跑但改动毫无效果，事后极难发现。
    """
    for item in overrides:
        key, value = parse_override(item)
        parts = key.split(".")
        node = cfg
        for depth, part in enumerate(parts[:-1]):
            nxt = node.get(part)
            if not isinstance(nxt, dict):
                if strict and part not in node:
                    raise KeyError(_unknown_key_message(key, ".".join(parts[:depth]), node, part))
                nxt = {}
                node[part] = nxt
            node = nxt
        if strict and parts[-1] not in node:
            raise KeyError(_unknown_key_message(key, ".".join(parts[:-1]), node, parts[-1]))
        node[parts[-1]] = value
    return cfg


def _unknown_key_message(full_key: str, prefix: str, node: Mapping, missing: str) -> str:
    """报错时给出同级可用键与最接近的拼写，方便一眼看出打错在哪。"""
    import difflib

    available = sorted(k for k in node if not k.startswith("_"))
    close = difflib.get_close_matches(missing, available, n=3, cutoff=0.6)
    where = f"`{prefix}` 下" if prefix else "配置顶层"
    hint = f"，是不是想写 {close[0]!r}？" if close else ""
    return (
        f"--set 里的 {full_key!r} 在配置中不存在：{where}没有 {missing!r}{hint} "
        f"同级可用的键有 {available}。"
        "（只允许覆盖已存在的键；确实要新增请先写进 yaml）"
    )


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)