#!/usr/bin/python

#
# graphviz.py - a simple interface to the graphviz suite of programs
#
# copyright (c) 2003 - Patrick Wagstrom
#                      wagspat@iit.edu
#
# released under the terms of the GNU General Public License version 2
#

import popen2
nodeCounter = 0
graphCounter = 0
# perl's method for white space is a whole lot nicer...

class nodeEdgeBase:
    def __init__(self):
        self.props = {}
        
    def addProp(self, propname, propvalue):
        self.props[propname] = propvalue

    def delProp(self, propname):
        if self.props.has_key(propname):
            del self.props[propname]

    def __repr__(self):
        return self.serialize()
    
    def serialize(self):
        output = []
        nameLen = len(self.name)
        numProps = len(self.props.keys())
        ctr = 0
        if numProps == 0:
            return "%s;" % (self.name)
        ts = "%s [" % self.name
        for thisProp in self.props.keys():
            if ctr == 0:
                ts = "%s [" % self.name
            else:
                ts = " " * (nameLen+2)
            ts = ts + "%s=\"%s\"" % (thisProp, self.props[thisProp])
            if ctr + 1 < numProps:
                ts = ts + ","
            else:
                ts = ts + "];"
            output.append(ts)
            ctr = ctr + 1
        return "\n".join(output)
    
class node(nodeEdgeBase):
    def __init__(self):
        nodeEdgeBase.__init__(self)
        global nodeCounter
        self.name = "n%s" % (nodeCounter)
        nodeCounter = nodeCounter + 1

class edge(nodeEdgeBase):
    def __init__(self, *args):
        nodeEdgeBase.__init__(self)
        # self.props = {}
        self.nodes = []
        # if this is a digraph this should be "->"
        self.symbol = "--"
        for thisNode in args:
            self.nodes.append(thisNode)

    def addNode(self, node):
        if not node in self.nodes:
            self.nodes.append(node)
            
    def delNode(self, node):
        if node in self.nodes:
            self.nodes.remove(node)
            
    def serialize(self):
        self.name = (" %s " % self.symbol).join([x.name for x in self.nodes])
        return nodeEdgeBase.serialize(self)
        
class nodeDefaults(nodeEdgeBase):
    def __init__(self):
        nodeEdgeBase.__init__(self)
        self.name = "node"

    def serialize(self):
        if len(self.props.keys()) == 0:
            return ""
        return nodeEdgeBase.serialize(self)

class edgeDefaults(nodeEdgeBase):
    def __init__(self):
        nodeEdgeBase.__init__(self)
        self.name = "edge"

    def serialize(self):
        if len(self.props.keys()) == 0:
            return ""
        return nodeEdgeBase.serialize(self)
            
class graph:
    def __init__(self):
        global graphCounter
        self.edges = []
        self.nodes = []
        self.props = {}
        self.edgeDefaults = edgeDefaults()
        self.nodeDefaults = nodeDefaults()
        self.name = "g%s" % (graphCounter)
        graphCounter = graphCounter + 1
        # if I ever need to digraphs/subgraphs, change or subclass this...
        self.graphType = "graph"

    def addNodeDefault(self, prop, propvalue):
        self.nodeDefaults.addProp(prop,propvalue)
    def delNodeDefault(self, prop):
        self.nodeDefaults.delProp(prop)

    def addEdgeDefault(self, prop, propvalue):
        self.edgeDefaults.addProp(prop, propvalue)
    def delEdgeDefault(self, prop):
        self.edgeDefaults.delProp(prop)
        
    def addProp(self, propname, propvalue):
        self.props[propname] = propvalue

    def delProp(self, propname):
        if self.props.has_key(propname):
            del self.props[propname]
            
    def addNode(self, *args):
        for node in args:
            if not node in self.nodes:
                self.nodes.append(node)

    def delNode(self, *args):
        for node in args:
            if node in self.nodes:
                a.remove(node)

    def addEdge(self, *args):
        for edge in args:
            if not edge in self.edges:
                self.edges.append(edge)

    def delEdge(self, *args):
        for edge in args:
            if edge in self.edges:
                a.remove(edge)

    def __repr__(self):
        return self.serialize()
        
    def serialize(self):
        indent = 4
        output = []
        output.append("%s %s {" % (self.graphType, self.name))
        for thisProp in self.props.keys():
            ts = " " * indent 
            ts = ts + "%s=\"%s\";" % (thisProp, self.props[thisProp])
            output.append(ts)

        # serialize the defaults
        output = output + [" " * indent + x
                           for x in self.nodeDefaults.serialize().splitlines()]
        output = output + [" " * indent + x
                           for x in self.edgeDefaults.serialize().splitlines()]

        for thisNode in self.nodes:
            output = output + [" " * indent + x
                               for x in thisNode.serialize().splitlines()]
        for thisEdge in self.edges:
            output = output + [" " * indent + x
                               for x in thisEdge.serialize().splitlines()]
        output.append("}")
        return "\n".join(output)
    

def renderGraph(inGraph,format="png",engine="dot"):
    r,w,e = popen2.popen3("%s -T%s" % (engine, format))
    w.write(inGraph.serialize())
    w.close()
    body = "".join(r.readlines())
    r.close()
    e.close()
    return body

if __name__ == "__main__":
    g = graph()
    n1 = node()
    n2 = node()
    n1.addProp("color","blue")
    e = edge(n1,n2)
    g.addNode(n1,n2)
    g.addEdge(e)
    g.addProp("size","3,3")
    g.addProp("bgcolor","#f0f0d0")
    g.addEdgeDefault("color","red")
    print g.serialize()

    print "\n\nAs SVG:"
    print renderGraph(g,format="svg")

    print "\n\nPNG output written to test.png"
    png = renderGraph(g,format="png")
    f = open("test.png","wb")
    f.write(png)
    f.close()