yamlcf.py

import os.path

import yaml


class Loader(yaml.SafeLoader):
    pass

def construct_include(loader, node):
    '''Evaluate contents of another yaml file, optionally selecting part of it.
    
    !config!include FILE

    !config!include
      file: FILE
      selector: [foo, 3, quux]

    '''

    if hasattr(loader.stream, 'name'):
        basedir = os.path.dirname(loader.stream.name)

    else:
        basedir = '.'

    if node.id == 'scalar':
        fn, selector = loader.construct_scalar(node), None

    else:
        d = loader.construct_mapping(node)
        fn, selector = d.get('file'), d.get('selector')

    fn = os.path.join(basedir, fn)
    doc = loader.construct_object(yaml.compose(file(fn), Loader=type(loader)))
    if selector:
        return reduce(lambda d, k: d[k], selector, doc)
    return doc


Loader.add_constructor('yaml.inprivatepractice.com:config:include',
        construct_include)

def construct_extend(loader, node):
    '''Deep merge multiple mappings.

    The following:

        !config!extend
        - server:
            port: 3000
            workers: 4
        - server:
            port: 3011

    becomes:

        server:
            port: 3011
            workers: 4

    '''

    def dzip(items, k=None):
        if not items:
            return

        if isinstance(items[0], dict):
            dicts = [ a for a in items if isinstance(a, dict) ]
            keys = reduce(set.union, dicts, set())
            return {
                k: dzip([ d.get(k) for d in dicts if k in d ], k)
                for k in keys
            }

        else:
            return items[-1]

    return dzip(loader.construct_sequence(node, True))

Loader.add_constructor('yaml.inprivatepractice.com:config:extend',
        construct_extend)


def construct_multi_env(loader, suffix, node):
    '''Evaluate an environment variable.

    !env!PORT

    The environment variable is evaluated as a yaml document.  This means we get
    native types with yaml's relaxed quoting rules (e.g. "5112" is an
    integer, but "http://127.0.0.1:5012/" is a string).  It also means we can
    pass complex objects through the environment.  This could cause some
    compatability issues when mixed with software that does not use this
    convention.
    '''

    name = suffix+loader.construct_scalar(node)

    data = os.environ.get(name) or None

    if data:
        return yaml.load(data, Loader=type(loader))

Loader.add_multi_constructor('yaml.inprivatepractice.com:config:env:',
        construct_multi_env)


def test():
    from tempfile import NamedTemporaryFile
    from textwrap import dedent
    from pprint import pformat
    from mock import patch

    with NamedTemporaryFile() as include_target, \
        patch('os.environ', dict(
            PORT='5112',
            FOO_URL='http://127.0.0.1:5112')):

        include_target.write(dedent('''\
                server:
                    port: 3000
                    workers: 3000
                database: sqlite:///app.db
                '''))

        include_target.flush()

        config = yaml.load(dedent('''\
                %TAG !config! yaml.inprivatepractice.com:config:
                %TAG !env! yaml.inprivatepractice.com:config:env:
                %TAG !myenv! yaml.inprivatepractice.com:config:env:FOO_
                ---
                !config!extend
                - !config!include '''+include_target.name+'''
                - server:
                    port: !env!PORT
                    url: !myenv!URL
                  database: postgres:///foo
                '''), Loader=Loader)

        assert config == {
            'database': 'postgres:///foo',
            'server': {
                'port': 5112,
                'workers': 3000,
                'url': 'http://127.0.0.1:5112',
            }}, pformat(config, width=10)


if __name__ == '__main__':
    test()