pvl/dot.py
author Tero Marttila <terom@paivola.fi>
Mon, 09 Mar 2015 23:31:13 +0200
changeset 738 3104fdf7ea26
parent 737 6ba76ac0bc72
permissions -rw-r--r--
pvl.hosts.hosts: drop support for instanced ip.* in favor of improved interface:ip.* =
import contextlib
import string

"""
    Syntax sugar for generating .dot files
    
        http://www.graphviz.org/content/dot-language
        http://www.graphviz.org/content/attrs
"""

ID = string.ascii_letters + string.digits + '_'

def quote (value):
    """
        (quoted) dot word.

        >>> print quote('foo')
        foo
        >>> print quote(5)
        5
        >>> print quote("five two")
        "five two"
    """

    value = str(value)
    
    if all(c in ID for c in value):
        return value
    else:
        return '"{value}"'.format(value=value)

def attr_list (**attrs):
    for name, value in attrs.iteritems():
        if value is None:
            continue
        
        yield '{name}={value}'.format(name=name, value=quote(value))

def print_stmt (*line, **attrs) :
    """
        >>> print_stmt(quote("Foo"), foo="bar")
          Foo [foo=bar];
        >>> print_stmt(quote("Foo"), asdf=5, quux="Test Me")
          Foo [quux="Test Me", asdf=5];
    """

    print "  {line}{attrs};".format(
            line    = ' '.join(item for item in line if item is not None),
            attrs   = ' [' + ', '.join(attr_list(**attrs)) + ']' if attrs else '',
    )

def print_edge(left, edgeop, right, **attrs):
    print_stmt(quote(left), edgeop, quote(right), **attrs)

@contextlib.contextmanager
def print_block (*block, **attrs):
    """
        >>> with print_block('graph', 'foo'):
        ...     print_stmt('foo', '--', 'bar')
        graph foo {
          foo -- bar;
        }
    """

    print ' '.join(item for item in block if item is not None) + ' {'
    for attr in attr_list(**attrs):
        print '  {attr};'.format(attr=attr)
    yield
    print '}'

def print_node(node, **attrs):
    print_stmt(quote(node), **attrs)

def print_graph(id=None, **attrs):
    """
        >>> with print_graph(): print_node("foo", label="Foo Bar")
        graph {
          foo [label="Foo Bar"];
        }
    """

    return print_block('graph', quote(id) if id else None, **attrs)

def print_subgraph(id=None):
    return print_block('subgraph', quote(id) if id else None)