"""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] = _strip_replace_key(value) elif key in merged and isinstance(merged[key], Mapping) and isinstance(value, Mapping): merged[key] = deep_merge(merged[key], value) else: merged[key] = _strip_replace_key(value) if isinstance(value, Mapping) else value return merged def _strip_replace_key(value: Mapping) -> dict: """去掉 `_replace_` 标记本身,避免它漏进组件的构造参数。""" return { k: (_strip_replace_key(v) if isinstance(v, Mapping) else v) for k, v in value.items() if k != REPLACE_KEY } 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.setdefault("config_path", str(Path(path).resolve())) return Config(cfg)