Creating Custom Tag in PyYAML

Question

I'm trying to use Python's PyYAML to create a custom tag that will allow me to retrieve environment variables with my 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)

And inside of defaults.yaml is simply:

example: !ENV foo

The error I keep getting:

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

I plan to have more than one custom tag as well (assuming I can get this one working)

Answer 1

If your goal is to find and replace environment variables (as strings) defined in your yaml file, you can use the following approach:

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'])

It's worth noting, you don't necessarily need to create a custom EnvLoader class. You can call add_constructor directly on the SafeLoader class or the yaml module itself. However, this can have an unintended side-effect of adding your loader globally to all other modules that rely on those loaders, which could potentially cuase problems if those other modules have their own custom logic for loading that !ENV tag.

Answer 2

I'd like to share how I resolved this as an addendum to the great answers above provided by Anthon and Fredrick Brennan. Thank you for your help.

In my opinion, the PyYAML document isn't real clear as to when you might want to add a constructor via a class (or "metaclass magic" as described in the doc), which may involve re-defining from_yaml and to_yaml, or simply adding a constructor using yaml.add_constructor.

In fact, the doc states:

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

I would argue the opposite is true for simpler use-cases. Here's how I managed to implement my custom tag.

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]

With this, I can now have a seperate yaml for each environment using an env PTYHON_ENV (default.yaml, development.yaml, test.yaml, production.yaml). And each can now reference ENV variables.

Example default.yaml:

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

Example production.yaml:

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

To use:

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 3

Your PyYAML class had a few problems:

1. yaml_tag is case sensitive, so !Env and !ENV are different tags.
2. So, as per the documentation, yaml.YAMLObject uses meta-classes to define itself, and has default to_yaml and from_yaml functions for those cases. By default, however, those functions require that your argument to your custom tag (in this case !ENV) be a mapping. So, to work with the default functions, your defaults.yaml file must look like this (just for example) instead:

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

Your code will then work unchanged, in my case print(settings) now results in {'example': /home/Fred} But you're using load instead of safe_load -- in their answer below, Anthon pointed out that this is dangerous because the parsed YAML can overwrite/read data anywhere on the disk.

You can still easily use your YAML file format, example: !ENV foo—you just have to define an appropriate to_yaml and from_yaml in class EnvTag, ones that can parse and emit scalar variables like the string "foo".

So:

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)

When this program is run, it outputs:

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

This code has the benefit of (1) using the original pyyaml, so nothing extra to install and (2) adding a representer. :)

Answer 4

There are several problems with your code:

• !Env in your YAML file is not the same as !ENV in your code.
• You are missing the classmethod from_yaml that has to be provided for EnvTag.
• Your YAML document specifies a scalar for !Env, but the subclassing mechanism for yaml.YAMLObject calls construct_yaml_object which in turn calls construct_mapping so a scalar is not allowed.
• You are using .load(). This is unsafe, unless you have complete control over the YAML input, now and in the future. Unsafe in the sense that uncontrolled YAML can e.g. wipe or upload any information from your disc. PyYAML doesn't warn you for that possible loss.
• PyYAML only supports most of YAML 1.1, the latest YAML specification is 1.2 (from 2009).
• You should consistently indent your code at 4 spaces at every level (or 3 spaces, but not 4 at the first and 3 a the next level).
• your __repr__ doesn't return a string if the environment variable is not set, which will throw an error.

So change your code to:

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)

which gives:

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

Please note that I am using ruamel.yaml (disclaimer: I am the author of that package), so you can use YAML 1.2 (or 1.1) in your YAML file. With minor changes you can do the above with the old PyYAML as well.

You can do this by subclassing of YAMLObject as well, and in a safe way:

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)

This will give you the same results as above.