terom@737: import contextlib
terom@737: import string
terom@737: 
terom@737: """
terom@737:     Syntax sugar for generating .dot files
terom@737:     
terom@737:         http://www.graphviz.org/content/dot-language
terom@737:         http://www.graphviz.org/content/attrs
terom@737: """
terom@737: 
terom@737: ID = string.ascii_letters + string.digits + '_'
terom@737: 
terom@737: def quote (value):
terom@737:     """
terom@737:         (quoted) dot word.
terom@737: 
terom@737:         >>> print quote('foo')
terom@737:         foo
terom@737:         >>> print quote(5)
terom@737:         5
terom@737:         >>> print quote("five two")
terom@737:         "five two"
terom@737:     """
terom@737: 
terom@737:     value = str(value)
terom@737:     
terom@737:     if all(c in ID for c in value):
terom@737:         return value
terom@737:     else:
terom@737:         return '"{value}"'.format(value=value)
terom@737: 
terom@737: def attr_list (**attrs):
terom@737:     for name, value in attrs.iteritems():
terom@737:         if value is None:
terom@737:             continue
terom@737:         
terom@737:         yield '{name}={value}'.format(name=name, value=quote(value))
terom@737: 
terom@737: def print_stmt (*line, **attrs) :
terom@737:     """
terom@737:         >>> print_stmt(quote("Foo"), foo="bar")
terom@737:           Foo [foo=bar];
terom@737:         >>> print_stmt(quote("Foo"), asdf=5, quux="Test Me")
terom@737:           Foo [quux="Test Me", asdf=5];
terom@737:     """
terom@737: 
terom@737:     print "  {line}{attrs};".format(
terom@737:             line    = ' '.join(item for item in line if item is not None),
terom@737:             attrs   = ' [' + ', '.join(attr_list(**attrs)) + ']' if attrs else '',
terom@737:     )
terom@737: 
terom@737: def print_edge(left, edgeop, right, **attrs):
terom@737:     print_stmt(quote(left), edgeop, quote(right), **attrs)
terom@737: 
terom@737: @contextlib.contextmanager
terom@737: def print_block (*block, **attrs):
terom@737:     """
terom@737:         >>> with print_block('graph', 'foo'):
terom@737:         ...     print_stmt('foo', '--', 'bar')
terom@737:         graph foo {
terom@737:           foo -- bar;
terom@737:         }
terom@737:     """
terom@737: 
terom@737:     print ' '.join(item for item in block if item is not None) + ' {'
terom@737:     for attr in attr_list(**attrs):
terom@737:         print '  {attr};'.format(attr=attr)
terom@737:     yield
terom@737:     print '}'
terom@737: 
terom@737: def print_node(node, **attrs):
terom@737:     print_stmt(quote(node), **attrs)
terom@737: 
terom@737: def print_graph(id=None, **attrs):
terom@737:     """
terom@737:         >>> with print_graph(): print_node("foo", label="Foo Bar")
terom@737:         graph {
terom@737:           foo [label="Foo Bar"];
terom@737:         }
terom@737:     """
terom@737: 
terom@737:     return print_block('graph', quote(id) if id else None, **attrs)
terom@737: 
terom@737: def print_subgraph(id=None):
terom@737:     return print_block('subgraph', quote(id) if id else None)