place_bundle.py

#!/usr/bin/env python

import argparse
import pprint
import networkx as nx
import jujuclient
import yaml


def setup():
    parser = argparse.ArgumentParser()
    parser.add_argument('-a', '--annotate', action="store_true")
    parser.add_argument('-e', '--environment')
    parser.add_argument('-l', '--layout', default="circular")
    parser.add_argument('-o', '--output')
    parser.add_argument('-s', '--scale', type=float, default=1000.0)
    parser.add_argument('-v', '--view', action="store_true")
    parser.add_argument('bundle', type=file)

    options = parser.parse_args()

    algo = getattr(nx, options.layout + '_layout', None)
    if not algo:
        raise SystemExit('Unknown layout algorithm, try again, try harder.')
    options.algo = algo
    return options


def parse_bundle(fp):
    data = yaml.load(fp)
    data = data[data.keys()[0]]
    return data


def split_rel(r):
    return r.split(':', 1)


def mapply(func, g, **kwargs):
    args = {}
    for name in func.func_code.co_varnames:
        if name in kwargs:
            args[name] = kwargs[name]
    return func(g, **args)


def layout(bundle, algo, scale, options):
    g = nx.MultiGraph()
    for service in bundle['services']:
        g.add_node(service)

    for relation in bundle['relations']:
        src = split_rel(relation[0])[0]
        for tgt in relation[1]:
            tgt = split_rel(tgt)[0]
            g.add_edge(src, tgt)
    pos = mapply(algo, g, k=45, iterations=100)
    pprint.pprint(pos)

    #if options.view:
        #nx.draw(g)
        #plt.savefig('/tmp/bundle.png')
        #subprocess.call(['eog', '/tmp/bundle.png'])


    for service, data in bundle['services'].items():
        data['annotations'] = {
            "gui-x": float(pos[service][0]) * scale,
            "gui-y": float(pos[service][1]) * scale,
        }
        if options.annotate:
            options.env.set_annotation(service, 'service',
                                    data['annotations'])
            #print options.env.get_annotation(service, 'service')
    return g


def render_bundle(bundle, target):
    bundle = {'cloudfoundry': bundle}
    with open(target, 'w') as fp:
        yaml.dump(bundle, fp, default_flow_style=False)


def main():
    options = setup()
    if options.annotate:
        options.env = jujuclient.Environment.connect(options.environment)
    bundle = parse_bundle(options.bundle)
    layout(bundle, options.algo, options.scale, options)
    if options.output:
        render_bundle(bundle, options.output)


if __name__ == '__main__':
    main()