public
Description: Oliver's Python utilities
Homepage:
Clone URL: git://github.com/osteele/python-utils.git
python-utils / graphviz.py
100644 115 lines (104 sloc) 4.2 kb
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
"""An OO interface to the "GraphViz":http://www.research.att.com/sw/tools/graphviz/ open source graph drawing software.
"""
 
__author__ = "Oliver Steele <steele@osteele.com>"
__copyright__ = "Copyright 1999-2001 by Oliver Steele."
__license__ = "Python License"
__version__ = '1.0a2'
 
class Graph:
    def __init__(self, name='graph'):
        assert ' ' not in name
        self.name = name
        self.nodes = []
        self.nodeMap = {}
        self.edgeMap = {}
        self.nextId = 0
 
    def makeNextId(self, prefix='g'):
        self.nextId += 1
        return prefix + str(self.nextId)
 
    def addNode(self, name, **attrs):
        self.nodeMap.setdefault(name, {}).update(attrs)
        if name not in self.nodes:
            self.nodes.append(name)
 
    def addEdge(self, source, target, **attrs):
        self.addNode(source)
        self.addNode(target)
        self.edgeMap.setdefault((source, target), {}).update(attrs)
 
    def asSource(self):
        """Returns a string that can be printed by the DOT tool at
http://www.research.att.com/sw/tools/graphviz/ ."""
        import string
        output = []
        output.append('digraph %s {' % self.name);
        def attrString(attrs):
            if attrs:
                def a(k, v):
                    if k in ['URL', 'label']:
                        return '"%s"' % v
                    else:
                        return str(v)
                return ' [' + ', '.join(['%s=%s' % (k, a(k, v)) for k, v in attrs.items()]) + ']'
            return ''
        def ql(s): return '\"' + s + '\"'
        ranks = {}
        for name in self.nodes:
            attrs = {}
            attrs.update(self.nodeMap[name])
            if attrs.has_key('rank'):
                import sys
                ranks.setdefault(attrs['rank'], []).append(name)
                del attrs['rank']
            output.append('\t%s%s;' % (ql(name), attrString(attrs)))
        keys = ranks.keys()
        keys.sort()
        for rank in keys:
            nodes = ranks[rank]
            output.append('\t{rank=same; ' + ''.join(['%s; ' % ql(name) for name in nodes]) + '};')
        for (s0, s1), attrs in self.edgeMap.items():
            output.append('\t%s -> %s%s;' % (ql(s0), ql(s1), attrString(attrs)))
        output.append('}');
        return string.join(output, '\n')
 
    def write(self, fname, format):
        # format in text, ps, hpgl, pcl, ..., gif, jpeg, png, ismap, imap, svg, plain
        if format == 'dot':
            open(fname, 'w').write(self.asSource())
            return
        import os, tempfile
        dotter = 'dot'
        str = self.asSource()
        dotfile = tempfile.mktemp()
        try:
            open(dotfile, 'w').write(str)
            expr = ("%s -T%s %s -o \"%s\"" % (dotter, format, dotfile, fname))
            os.system(expr)
        finally:
            try: os.remove(dotfile)
            except: pass
 
    def as(self, format):
        if format == 'dot':
            return self.asSource()
        import tempfile
        fname = tempfile.mktemp()
        #fname = 'c:/documents and settings/steele/my documents/temp.out'
        self.write(fname, format)
        try:
            self.write(fname, format)
            return open(fname, 'rb').read()
        finally:
            try: os.remove(fname)
            except: pass
 
def _test():
    graph = Graph('my_graph')
    graph.addNode('Zeus', URL="zeus", shape='box', style='filled', rank=1)
    graph.addNode('Hera', URL="hera", rank=1)
    graph.addNode('Zeus+Hera', rank=1, shape='none')
    graph.addNode('Heracles', shape="box", rank=2)
    graph.addEdge('Zeus', 'Zeus+Hera', dir='none')
    graph.addEdge('Zeus+Hera', 'Hera', dir='none')
    graph.addEdge('Zeus+Hera', 'Heracles')
    #graph.addEdge('Zeus', 'Hera', label='wife', dir='none')
    #graph.addEdge('Zeus', 'Heracles', label='children')
    print graph.asSource()
    import os
    #graph.write(os.path.join(os.path.split(__file__)[0], 'zeus.dot'), 'dot')
    graph.write(os.path.join(os.path.split(__file__)[0], 'zeus.jpeg'), 'jpeg')
 
#_test()