在 PyYAML 中创建自定义标签

Question

我正在尝试使用 Python 的 PyYAML 创建一个自定义标签，允许我使用我的 YAML 检索环境变量。

import os
import yaml

class EnvTag(yaml.YAMLObject):
    yaml_tag = u'!Env'

    def __init__(self, env_var):
       self.env_var = env_var

    def __repr__(self):
       return os.environ.get(self.env_var)

settings_file = open('conf/defaults.yaml', 'r')
settings = yaml.load(settings_file)

而 defaults.yaml 的内部就是：

example: !ENV foo

我不断收到的错误：

yaml.constructor.ConstructorError: 
could not determine a constructor for the tag '!ENV' in 
"defaults.yaml", line 1, column 10

我也计划拥有不止一个自定义标签（假设我可以让这个正常工作）

Answer 1

你的 PyYAML class 有一些问题：

yaml_tag区分大小写，所以!Env和!ENV是不同的标签。
因此，根据文档，yaml.YAMLObject 使用 meta-classes 来定义自身，并且在这些情况下具有默认的 to_yaml 和 from_yaml 函数。但是，默认情况下，这些函数要求您对自定义标记（在本例中为 !ENV）的参数是映射。因此，要使用默认函数，您的 defaults.yaml 文件必须如下所示（仅作为示例）：

example: !ENV {env_var: "PWD", test: "test"}

然后您的代码将保持不变，在我的例子中 print(settings) 现在导致 {'example': /home/Fred} 但是您使用的是 load 而不是 safe_load -- 在下面的回答中，Anthon 指出这很危险，因为解析后的 YAML 可以 overwrite/read 磁盘上任何位置的数据。

您仍然可以轻松使用您的 YAML 文件格式，example: !ENV foo——您只需在 class EnvTag 中定义适当的 to_yaml 和 from_yaml ，可以解析和发出标量变量的变量，例如字符串 "foo".

所以：

import os
import yaml

class EnvTag(yaml.YAMLObject):
    yaml_tag = u'!ENV'

    def __init__(self, env_var):
        self.env_var = env_var

    def __repr__(self):
        v = os.environ.get(self.env_var) or ''
        return 'EnvTag({}, contains={})'.format(self.env_var, v)

    @classmethod
    def from_yaml(cls, loader, node):
        return EnvTag(node.value)

    @classmethod
    def to_yaml(cls, dumper, data):
        return dumper.represent_scalar(cls.yaml_tag, data.env_var)

# Required for safe_load
yaml.SafeLoader.add_constructor('!ENV', EnvTag.from_yaml)
# Required for safe_dump
yaml.SafeDumper.add_multi_representer(EnvTag, EnvTag.to_yaml)

settings_file = open('defaults.yaml', 'r')

settings = yaml.safe_load(settings_file)
print(settings)

s = yaml.safe_dump(settings)
print(s)

当这个程序是运行时，输出：

{'example': EnvTag(foo, contains=)}
{example: !ENV 'foo'}

此代码的好处是 (1) 使用原始的 pyyaml，因此无需额外安装，以及 (2) 添加代表。 :)

Answer 2

您的代码有几个问题：

!Env 与代码中的 !ENV 不同。
您缺少必须为 EnvTag 提供的 classmethod from_yaml。
您的 YAML 文档为 !Env 指定了一个标量，但是 yaml.YAMLObject 的子类化机制调用了 construct_yaml_object，后者又调用了 construct_mapping，因此不允许使用标量。
您正在使用 .load()。这是 不安全，除非您现在和将来都可以完全控制 YAML 输入。从某种意义上说，不受控制的 YAML 可能是不安全的，例如从您的光盘擦除或上传任何信息。 PyYAML 不会警告您可能的损失。
PyYAML 只支持大部分 YAML 1.1，最新的 YAML 规范是 1.2（从 2009 年开始）。
您应该始终如一地在每个级别将代码缩进 4 个空格（或 3 个空格，但不是第一级为 4 个，下一级为 3 个）。
如果未设置环境变量，您的 __repr__ 不会 return 字符串，这将引发错误。

因此将您的代码更改为：

import sys
import os
from ruamel import yaml

yaml_str = """\
example: !Env foo
"""


class EnvTag:
    yaml_tag = u'!Env'

    def __init__(self, env_var):
        self.env_var = env_var

    def __repr__(self):
        return os.environ.get(self.env_var, '')

    @staticmethod
    def yaml_constructor(loader, node):
        return EnvTag(loader.construct_scalar(node))


yaml.add_constructor(EnvTag.yaml_tag, EnvTag.yaml_constructor,
                     constructor=yaml.SafeConstructor)

data = yaml.safe_load(yaml_str)
print(data)
os.environ['foo'] = 'Hello world!'
print(data)

给出：

{'example': }
{'example': Hello world!}

请注意，我使用的是 ruamel.yaml（免责声明：我是该软件包的作者），因此您可以在 YAML 文件中使用 YAML 1.2（或 1.1）。稍作改动，您也可以使用旧的 PyYAML 执行上述操作。

您也可以通过 YAMLObject 的子类化来做到这一点，而且是安全的方式：

import sys
import os
from ruamel import yaml

yaml_str = """\
example: !Env foo
"""

yaml.YAMLObject.yaml_constructor = yaml.SafeConstructor

class EnvTag(yaml.YAMLObject):
    yaml_tag = u'!Env'

    def __init__(self, env_var):
        self.env_var = env_var

    def __repr__(self):
        return os.environ.get(self.env_var, '')

    @classmethod
    def from_yaml(cls, loader, node):
        return EnvTag(loader.construct_scalar(node))


data = yaml.safe_load(yaml_str)
print(data)
os.environ['foo'] = 'Hello world!'
print(data)

这将为您提供与上述相同的结果。

Answer 3

我想分享一下我是如何解决这个问题的，作为 Anthon 和 Fredrick Brennan 提供的上述出色答案的补充。感谢您的帮助。

在我看来，PyYAML 文档并不清楚您何时可能希望通过 class（或文档中描述的 "metaclass magic" 添加构造函数），这可能涉及重新定义 from_yaml 和 to_yaml，或者简单地使用 yaml.add_constructor 添加构造函数。

事实上，doc 声明：

You may define your own application-specific tags. The easiest way to do it is to define a subclass of yaml.YAMLObject

我认为对于更简单的用例情况恰恰相反。以下是我如何成功实施我的自定义标签。

config/__init__.py

import yaml
import os

environment = os.environ.get('PYTHON_ENV', 'development')

def __env_constructor(loader, node):
    value = loader.construct_scalar(node)
    return os.environ.get(value)

yaml.add_constructor(u'!ENV', __env_constructor)

# Load and Parse Config
__defaults      = open('config/defaults.yaml', 'r').read()
__env_config    = open('config/%s.yaml' % environment, 'r').read()
__yaml_contents = ''.join([__defaults, __env_config])
__parsed_yaml   = yaml.safe_load(__yaml_contents)

settings = __parsed_yaml[environment]

有了这个，我现在可以使用 env PTYHON_ENV (default.yaml, development.yaml, test.yaml, production.yaml 为每个环境创建一个单独的 yaml ).每个人现在都可以引用 ENV 变量。

示例default.yaml：

defaults: &default
  app:
    host: '0.0.0.0'
    port: 500

示例production.yaml：

production:
  <<: *defaults
  app:
    host: !ENV APP_HOST
    port: !ENV APP_PORT

使用：

from config import settings
"""
If PYTHON_ENV == 'production', prints value of APP_PORT
If PYTHON_ENV != 'production', prints default 5000
"""
print(settings['app']['port'])

Answer 4

如果您的目标是查找和替换 yaml 文件中定义的环境变量（作为字符串），您可以使用以下方法：

example.yaml:

foo: !ENV "Some string with ${VAR1} and ${VAR2}"

example.py:

import yaml

# Define the function that replaces your env vars
def env_var_replacement(loader, node):
    replacements = {
      '${VAR1}': 'foo',
      '${VAR2}': 'bar',
    }
    s = node.value
    for k, v in replacements.items():
        s = s.replace(k, v)
    return s

# Define a loader class that will contain your custom logic
class EnvLoader(yaml.SafeLoader):
    pass

# Add the tag to your loader
EnvLoader.add_constructor('!ENV', env_var_replacement)

# Now, use your custom loader to load the file:
with open('example.yaml') as yaml_file:
    loaded_dict = yaml.load(yaml_file, Loader=EnvLoader)

    # Prints: "Some string with foo and bar"
    print(loaded_dict['foo'])

值得注意的是，您不一定需要创建自定义 EnvLoader class。您可以直接在 SafeLoader class 或 yaml 模块本身上调用 add_constructor。然而，这可能会产生一个意想不到的副作用，将你的加载器全局添加到依赖这些加载器的所有其他模块，如果其他模块有自己的自定义逻辑来加载 !ENV 标签，这可能会导致问题。

在 PyYAML 中创建自定义标签

Creating Custom Tag in PyYAML

python

tags

pyyaml