1

import contextlib

     2

import string

     3

     4

"""

     5

    Syntax sugar for generating .dot files

     6

     7

        http://www.graphviz.org/content/dot-language

     8

        http://www.graphviz.org/content/attrs

     9

"""

    10

    11

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

    12

    13

def quote (value):

    14

    """

    15

        (quoted) dot word.

    16

    17

        >>> print quote('foo')

    18

        foo

    19

        >>> print quote(5)

    20

        5

    21

        >>> print quote("five two")

    22

        "five two"

    23

    """

    24

    25

    value = str(value)

    26

    27

    if all(c in ID for c in value):

    28

        return value

    29

    else:

    30

        return '"{value}"'.format(value=value)

    31

    32

def attr_list (**attrs):

    33

    for name, value in attrs.iteritems():

    34

        if value is None:

    35

            continue

    36

    37

        yield '{name}={value}'.format(name=name, value=quote(value))

    38

    39

def print_stmt (*line, **attrs) :

    40

    """

    41

        >>> print_stmt(quote("Foo"), foo="bar")

    42

          Foo [foo=bar];

    43

        >>> print_stmt(quote("Foo"), asdf=5, quux="Test Me")

    44

          Foo [quux="Test Me", asdf=5];

    45

    """

    46

    47

    print "  {line}{attrs};".format(

    48

            line    = ' '.join(item for item in line if item is not None),

    49

            attrs   = ' [' + ', '.join(attr_list(**attrs)) + ']' if attrs else '',

    50

    )

    51

    52

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

    53

    print_stmt(quote(left), edgeop, quote(right), **attrs)

    54

    55

@contextlib.contextmanager

    56

def print_block (*block, **attrs):

    57

    """

    58

        >>> with print_block('graph', 'foo'):

    59

        ...     print_stmt('foo', '--', 'bar')

    60

        graph foo {

    61

          foo -- bar;

    62

        }

    63

    """

    64

    65

    print ' '.join(item for item in block if item is not None) + ' {'

    66

    for attr in attr_list(**attrs):

    67

        print '  {attr};'.format(attr=attr)

    68

    yield

    69

    print '}'

    70

    71

def print_node(node, **attrs):

    72

    print_stmt(quote(node), **attrs)

    73

    74

def print_graph(id=None, **attrs):

    75

    """

    76

        >>> with print_graph(): print_node("foo", label="Foo Bar")

    77

        graph {

    78

          foo [label="Foo Bar"];

    79

        }

    80

    """

    81

    82

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

    83

    84

def print_subgraph(id=None):

    85

    return print_block('subgraph', quote(id) if id else None)