如何使用ruamel.yaml为OmegaConf从dataclass生成的YAML添加注释

直接通过ruamel.yaml的带注释节点能力，搭配docstring注释解析即可实现需求，不需要复杂正则匹配，也能完美处理多层级嵌套的注释挂载问题。

前置依赖 #

需要先安装两个依赖库：

pip install omegaconf ruamel.yaml

完整代码示例 #

from dataclasses import dataclass
from omegaconf import OmegaConf
from ruamel.yaml import YAML
from ruamel.yaml.comments import CommentedMap
import sys

def parse_doc_comments(doc_str: str) -> dict:
    """解析dataclass的__doc__为 键路径:注释 的字典"""
    comments = {}
    if not doc_str:
        return comments
    for line in doc_str.strip().split("\n"):
        line = line.strip()
        if not line or ":" not in line:
            continue
        key_path, comment = line.split(":", 1)
        comments[key_path.strip()] = comment.strip()
    return comments

def add_comments_to_node(node: CommentedMap, key_path: str, comment: str):
    """按层级键路径给对应节点添加行尾注释"""
    keys = key_path.split(".")
    current = node
    # 遍历处理父级节点，不存在则自动创建
    for key in keys[:-1]:
        if key not in current:
            current[key] = CommentedMap()
        current = current[key]
    final_key = keys[-1]
    # 给最终节点加行尾注释
    current.yaml_add_eol_comment(comment, final_key, column=0)

def omegaconf_to_commented_yaml(conf, doc_str: str):
    """将OmegaConf配置转为带注释的YAML"""
    # 先把OmegaConf转成普通字典，再转成ruamel的CommentedMap
    conf_dict = OmegaConf.to_container(conf)
    root = CommentedMap(conf_dict)
    # 解析注释
    comment_map = parse_doc_comments(doc_str)
    # 挂载所有注释
    for key_path, comment in comment_map.items():
        add_comments_to_node(root, key_path, comment)
    # 输出YAML
    yaml = YAML()
    yaml.indent(mapping=4, sequence=4, offset=2)
    yaml.dump(root, sys.stdout)

# 示例dataclass
@dataclass
class MyConfig:
    """
    port: This is a comment
    host: This is another comment
    section.another: Commenting a value under a base section
    """
    port: int = 80
    host: str = "localhost"

if __name__ == "__main__":
    conf = OmegaConf.structured(MyConfig)
    omegaconf_to_commented_yaml(conf, MyConfig.__doc__)

输出结果 #

运行上述代码会直接输出符合要求的带注释YAML：

port: 80 # This is a comment
host: localhost # This is another comment
section:
    another: # Commenting a value under a base section

拓展说明 #

• 如果需要把注释放在键的上方而非行尾，将yaml_add_eol_comment替换为yaml_set_comment_before_after_key方法即可
• 键路径匹配严格区分大小写，需要和docstring里的写法完全对应
• 未在dataclass中定义的嵌套父节点会自动创建，无需提前声明

内容的提问来源于stack exchange，提问作者Marcel Alexandru