Creating Custom Tag in PyYAML

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. :)

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.

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