如何在无ruamel时加载带!Employee属性标签的Python YAML文件

问题分析 #

你需要解析的YAML结构中，!Employee标签直接跟在状态键（如old）后面形成!Employee old:的紧凑写法，但PyYAML默认会将!Employee标签应用在old这个标量键上，而非整个键值对。原有代码仅能处理标签单独一行的合法YAML结构，无法适配这种紧凑写法。

解决方案 #

我们通过自定义Loader，在解析根映射时识别带有!Employee标签的键，将对应的键值对转换为Employee对象；同时调整Representer，保证导出时输出!Employee old:的紧凑格式。

修改后的完整代码 #

import yaml


class Employee:
    def __init__(self, status, data=None):
        self.status = status  # 如"old"
        self.data = data or {}  # 如{'name': 'John', 'id': 'some_id'}

    def __repr__(self):
        return f"Employee(status='{self.status}', data={self.data})"


# 自定义构造函数：处理标签应用于ScalarNode（即!Employee old中的old）的情况
def employee_constructor(loader, node):
    # 当标签应用于ScalarNode时，这个标量就是status
    if isinstance(node, yaml.ScalarNode):
        status = loader.construct_scalar(node)
        return status
    # 兼容原有标签单独一行的情况（标签应用于MappingNode）
    elif isinstance(node, yaml.MappingNode):
        first_key_node, value_node = node.value[0]
        status = loader.construct_scalar(first_key_node)
        data = loader.construct_mapping(value_node)
        return Employee(status, data)
    raise yaml.constructor.ConstructorError(None, None, f"不支持的节点类型: {type(node)}", node.start_mark)


# 自定义Representer：输出!Employee old:的紧凑格式
def employee_representer(dumper, data):
    # 创建一个带标签的ScalarNode作为键，值为data.data
    tagged_key = dumper.represent_scalar('!Employee', data.status)
    # 返回一个仅包含该键值对的MappingNode
    return dumper.represent_mapping('', {tagged_key: data.data}, flow_style=False)


# 自定义Loader：重写映射构造，处理带!Employee标签的键
class EmployeeLoader(yaml.FullLoader):
    def construct_mapping(self, node, deep=True):
        mapping = super().construct_mapping(node, deep=deep)
        new_mapping = {}
        for key, value in mapping.items():
            # 检查键是否是被!Employee标签标记的status
            if isinstance(key, str) and hasattr(key, '_employee_tag'):
                # 将该键值对转换为Employee对象，用固定键名如'employee'存储
                new_mapping['employee'] = Employee(key, value)
            else:
                new_mapping[key] = value
        return new_mapping


# 给通过ScalarNode构造的status字符串添加标记
def wrapped_employee_constructor(loader, node):
    result = employee_constructor(loader, node)
    if isinstance(result, str):
        # 添加自定义标记，用于后续识别
        setattr(result, '_employee_tag', True)
    return result


# 注册构造器和代表器
yaml.add_constructor('!Employee', wrapped_employee_constructor, Loader=EmployeeLoader)
yaml.add_representer(Employee, employee_representer)


def load_yaml(file_path):
    with open(file_path, 'r') as file:
        return yaml.load(file, Loader=EmployeeLoader)


def save_yaml(file_path, data):
    with open(file_path, 'w') as file:
        yaml.dump(data, file, default_flow_style=False)


if __name__ == "__main__":
    yaml_string = """
    !Employee old:
        name: Some_Name
        id: some_id
    
    cfg:
        param : 'abc'
    """

    # 加载YAML数据
    data = yaml.load(yaml_string, Loader=EmployeeLoader)
    print("加载后的YAML数据:")
    print(data)

    # 修改数据
    data['cfg']['param'] = 'xyz'
    data['employee'].data['name'] = 'New_Name'

    # 导出为YAML
    output = yaml.dump(data, default_flow_style=False)
    print("\n修改后的YAML输出:")
    print(output)
    save_yaml("output3.yaml", data)

代码说明 #

1. 构造逻辑调整：

   • 新增wrapped_employee_constructor，给从ScalarNode构造的status字符串添加自定义标记_employee_tag，方便后续识别。
   • 自定义EmployeeLoader重写construct_mapping方法，遍历映射键值对，识别带标记的status键，将其与对应值转换为Employee对象，存入根映射的employee键下。
   • 保留原有对MappingNode的处理，兼容标签单独一行的情况。

2. 导出逻辑调整：

   • employee_representer创建带!Employee标签的标量作为键，值为员工数据，输出!Employee old:的紧凑格式。

3. 效果验证：

   • 加载!Employee old:格式的YAML后，根映射会包含employee键（对应Employee对象）和cfg键。
   • 修改数据后导出，会保持!Employee old:的紧凑写法，同时保留cfg的同级结构。

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