Fix tipples

tipples.py

#!/usr/bin/env python3

import argparse
import random
import string
import sys
import textwrap
import yaml

r = lambda: random.randint(0,255)



def node(s):
    allowed = string.ascii_lowercase
    return ''.join(filter(string.ascii_lowercase.__contains__, s.lower()))

def getwrapper(width):
    def wrapper(s):
        return '\\n'.join(textwrap.wrap(s, width=width, break_long_words=False))

    return wrapper

def main():
    parser = argparse.ArgumentParser(description="Generate dot file from some YAML")
    parser.add_argument('-w', '--width', help="Width of each node and edge label for wrapping (default: %(default)s)", type=int, default=10)
    parser.add_argument('-i', '--input', help="Input yaml file, defaults to stdin", type=argparse.FileType('r'), default=sys.stdin)
    parser.add_argument('-o', '--output', help="Output graphviz dot file", type=argparse.FileType('w'), default=sys.stdout)
    args = parser.parse_args()

    wrapper = getwrapper(args.width)

    with args.input as input:
        graph = yaml.load(input)

    # make sure drinks are array, even if there is just one
    for group in graph.values():
        for person, drinks in group.items():
            if isinstance(drinks, str):
                group[person] = [drinks]

    with args.output as output:
        output.write('digraph Tipples {\n')
        output.write('  label = "{}";\n'.format(wrapper('Famous Tipples')))

        for category, people in graph.items():
            output.write('  subgraph cluster_{name} {{\n'.format(name=node(category)))
            output.write('    label = "{label}";\n'.format(label=wrapper(category)))
            output.write('    color = "#{r:02X}{g:02X}{b:02X}";\n'.format(r=r(), g=r(), b=r()))

            for person in people.keys():
                output.write('    {node} [label = "{name}"];\n'.format(node=node(person), name=wrapper(person)))

            output.write('  }\n')

        drinks = frozenset(drink for group in graph.values() for drinks in group.values() for drink in drinks)

        for drink in drinks:
            output.write('  {node} [label = "{name}"];\n'.format(node=node(drink), name=wrapper(drink)))

        for group in graph.values():
            for person, drinks in group.items():
                for drink in drinks:
                    output.write('  {person} -> {drink};\n'.format(person=node(person), drink=node(drink)))
        output.write('}\n')



if __name__ == '__main__':
    main()