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

本文主要是介绍PyYAML中自定义tag标签获取环境变量值，对大家解决编程问题具有一定的参考价值，需要的程序猿们随着小编来一起学习吧！

前言

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

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

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

自定义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官方文档

这篇关于PyYAML中自定义tag标签获取环境变量值的文章就介绍到这儿，希望我们推荐的文章对大家有所帮助，也希望大家多多支持为之网！