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官方文檔