如何在结合Pydantic的YAML文件中复用变量

需求描述 #

我希望加载YAML文件并创建Pydantic BaseModel对象，同时实现YAML文件内部的变量复用。示例YAML文件如下：

config:
  variables:
    root_level: DEBUG
    my_var: "TEST"

  handlers_logs:
    - class: $my_var #<--- 需要替换的变量
      level_threshold: STATS
      block_level_filter: true
      disable: false
      args:
        hosts: $my_var #<--- 需要替换的变量
        topic: _stats

现有代码如下：

import os
from pprint import pprint

import yaml
from pydantic import BaseModel
from typing import Dict, Optional, Any
from yaml.parser import ParserError

class BaseLogModel(BaseModel):
    class Config:
        use_enum_values = True
        allow_population_by_field_name = True

class Config(BaseLogModel):
    variables: Optional[Dict[str, str]]
    handlers_logs: Any

def load_config(filename) -> Optional[Config]:
    if not os.path.exists(filename):
        return None

    with open(filename) as f:
        try:
            config_file = yaml.load(f.read(), Loader=yaml.SafeLoader)
            if config_file is not None and isinstance(config_file, dict):
                config_data = config_file["config"]
            else:
                return None
        except ParserError as e:
            return None

    return Config.parse_obj(config_data)

def main():
    config = load_config("config.yml")
    pprint(config)

当前输出：

Config(variables={'root_level': 'DEBUG', 'my_var': 'TEST'}, handlers_logs=[{'class': '$my_var', 'level_threshold': 'STATS', 'block_level_filter': True, 'disable': False, 'args': {'hosts': '$my_var', 'topic': '_stats'}}])

我希望将其中的$my_var替换为"TEST"，避免重复编写相同值，请问能否通过Pydantic或其他YAML库实现这一需求？

解决方案 #

方法1：加载YAML后手动递归替换变量 #

先加载原始YAML数据，提取variables中的变量映射，再递归遍历整个配置数据，将所有$变量名格式的字符串替换为对应的值，最后传给Pydantic解析。

修改后的load_config函数：

def replace_variables(data, variables):
    if isinstance(data, dict):
        return {k: replace_variables(v, variables) for k, v in data.items()}
    elif isinstance(data, list):
        return [replace_variables(item, variables) for item in data]
    elif isinstance(data, str) and data.startswith('$'):
        var_name = data[1:]
        return variables.get(var_name, data)  # 变量不存在则保留原字符串
    else:
        return data

def load_config(filename) -> Optional[Config]:
    if not os.path.exists(filename):
        return None

    with open(filename) as f:
        try:
            config_file = yaml.load(f.read(), Loader=yaml.SafeLoader)
            if config_file is not None and isinstance(config_file, dict):
                config_data = config_file["config"]
            else:
                return None
        except ParserError as e:
            return None

    # 提取变量并替换
    variables = config_data.get('variables', {})
    config_data = replace_variables(config_data, variables)

    return Config.parse_obj(config_data)

运行后输出：

Config(variables={'root_level': 'DEBUG', 'my_var': 'TEST'}, handlers_logs=[{'class': 'TEST', 'level_threshold': 'STATS', 'block_level_filter': True, 'disable': False, 'args': {'hosts': 'TEST', 'topic': '_stats'}}])

方法2：自定义YAML Loader实现变量替换 #

继承yaml.SafeLoader实现自定义Loader，在加载YAML时直接完成变量替换，无需后续处理。

示例代码：

import yaml
from yaml.nodes import ScalarNode

class VariableLoader(yaml.SafeLoader):
    def __init__(self, stream, variables):
        super().__init__(stream)
        self.variables = variables

    def construct_scalar(self, node):
        value = super().construct_scalar(node)
        if isinstance(value, str) and value.startswith('$'):
            var_name = value[1:]
            return self.variables.get(var_name, value)
        return value

def load_config(filename) -> Optional[Config]:
    if not os.path.exists(filename):
        return None

    with open(filename) as f:
        try:
            # 先加载一次获取variables
            raw_config = yaml.load(f.read(), Loader=yaml.SafeLoader)
            if not raw_config or not isinstance(raw_config, dict):
                return None
            variables = raw_config.get('config', {}).get('variables', {})
            
            # 重新打开文件，用自定义Loader加载并替换变量
            f.seek(0)
            config_file = yaml.load(f.read(), Loader=lambda stream: VariableLoader(stream, variables))
            if config_file is not None and isinstance(config_file, dict):
                config_data = config_file["config"]
            else:
                return None
        except ParserError as e:
            return None

    return Config.parse_obj(config_data)

方法3：利用Pydantic的validator在模型解析时替换 #

通过Pydantic的@validator装饰器，将变量替换逻辑绑定到模型上，在验证阶段自动完成替换。

修改模型类：

from pydantic import validator

class Config(BaseLogModel):
    variables: Optional[Dict[str, str]] = {}
    handlers_logs: Any

    @validator('handlers_logs')
    def replace_handlers_variables(cls, v, values):
        variables = values.get('variables', {})
        
        def _replace(data):
            if isinstance(data, dict):
                return {k: _replace(val) for k, val in data.items()}
            elif isinstance(data, list):
                return [_replace(item) for item in data]
            elif isinstance(data, str) and data.startswith('$'):
                var_name = data[1:]
                return variables.get(var_name, data)
            else:
                return data
        
        return _replace(v)

这种方法无需修改加载函数，Pydantic解析模型时会自动处理变量替换。

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