PyYAML中自定义tag标签获取环境变量值

前言

YAML常用于配置文件，当配置文件中需要配置一些用户名密码时，直接写在YAML文件并上传到代码仓库中则很容易造成密码泄露。
不幸的是，前一段时间我们组的自动化代码就被检测到了密码泄露，被通知整改。

yaml使用基础，参考：https://www.cnblogs.com/superhin/p/11503756.html

解决的方法有两种：

1. 配置文件仅本地使用，不传到代码仓库中
2. 将密码配置到执行机器的环境变量中，在YAML中使用特殊标记表示读取一个环境变量

我们可以使用自定义tag来实现这种功能。在PyYAML中一种tag标识一种类型，常见的tag有：

!!null	None
!!bool	bool
!!int	int
!!float	float
!!binary	bytes
!!timestamp	datetime.datetime
!!omap, !!pairs	list of pairs
!!set	set
!!str	str
!!seq	list
!!map	dict

自定义tag

我们自定义一个新的tag, !env, 并编写一个对应的处理函数（PyYAML中称为constructor构造器），代码如下：
demo.yaml文件

user: !env ${USER}  # 表示环境变量USER，即当前用户名

Python文件如下：

import os
import yaml  # 需要pip install pyyaml

def env_var_constructor(loader, node):
    value = loader.construct_scalar(node)   # PyYAML loader的固定方法，用于根据当前节点构造一个变量值
    var_name = value.strip('${} ')  # 去除变量值（例如${USER}）前后的特殊字符及空格
    return os.getenv(var_name, value)  # 尝试在环境变量中获取变量名（如USER）对应的值，获取不到使用默认值value（即原来的${USER}）

yaml.SafeLoader.add_constructor('!env', env_var_constructor)  # 为SafeLoader添加新的tag和构造器

with open('demo.yml') as f:  
    print(yaml.safe_load(f))  # 打开文件并使用SafeLoader加载文件内容

结果如下：

{'user': 'superhin'}

为tag分配匹配模式

此时YAML文件中环境变量只能使用强制类型声明!env ${变量名}来使用，如果想直接使用${变量名}来使用则需要为该tag指定一种正则匹配模式，即识别到类似${变量名}格式时自动使用!env这个tag。
demo.yaml文件

user: !env ${USER}  # 表示环境变量USER，即当前用户名
path: ${PATH}  # 期望可以直接使用

import os
import re
import yaml

pattern = re.compile('\${\w+}')  # 匹配 ${一个或多个字母或数字}

def env_var_constructor(loader, node):
    value = loader.construct_scalar(node)  
    var_name = value.strip('${} ')  
    return os.getenv(var_name, value)  

yaml.SafeLoader.add_constructor('!env', env_var_constructor)  # 添加新tag即对应的构造器
yaml.SafeLoader.add_implicit_resolver('!env', pattern, None)    #  为tag指定一种正则匹配

with open('demo.yml') as f:  
    print(yaml.safe_load(f))  # 打开文件并使用SafeLoader加载文件内容

运行结果如下

{'user': 'superhin', 'path': '...<省略>'}

一个节点使用多个变量

如果我们想要在一个节点中使用多个变量，如
demo.yml内容

user: !env ${USER}
path: ${PATH}
msg: 当前用户名 ${USER} 系统路径 ${PATH}

则需要对节点值value(字符串格式)进行逐个替换。
首先我们需要修改我们的匹配模式，允许${变量}前后可以拥有多个任意字符

pattern = re.compile('.*?(\${\w+}).*?')  # 前后可以拥有多个任意字符，使用小括号分组只取当前变量${变量名}内容，`?`表示非贪婪匹配。

完整代码如下：

import os
import re
import yaml

pattern = re.compile('.*?(\${\w+}).*?')  # 前后可以拥有多个任意字符，使用小括号分组只取当前变量${变量名}内容，`?`表示非贪婪匹配。

def env_var_constructor(loader, node):
    value = loader.construct_scalar(node)
    for item in pattern.findall(value):  # 遍历所有匹配到到${变量名}的变量, 如${USER}
        var_name = item.strip('${} ')    # 如,USER
        value = value.replace(item, os.getenv(var_name, item))  # 用环境变量中取到的对应值替换当前变量
    return value                                  # 如superin替换${USER}，取不到则使用原值${USER}

yaml.SafeLoader.add_constructor('!env', env_var_constructor)
yaml.SafeLoader.add_implicit_resolver('!env', pattern, None)

with open('demo.yml') as f:
    print(yaml.safe_load(f))

运行结果如下：

{'user': 'superhin', 'path': '...<省略>', 'msg': '当前用户名 superhin 系统路径 ...<省略>'}

参考：PyYAML官方文档