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 En         


        

Answer 1

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.