Using graphviz's "dot" to arrange a Nuke node graph

nuke_dot_arrange.py

"""Test of using graphviz's "dot" layout algorithm to arrange a Nuke node-graph
Example: http://i.imgur.com/f1mK7.png
Left is dot, right is Nuke's builtin nuke.autoplace
"""

import nuke

# http://code.google.com/p/pydot/
import pydot


class NukeDoter(object):
    def __init__(self):
        self.graph = pydot.Dot(graph_type='digraph')
        self.checked = []
        self.counter = 0

    def recursive_grapher(self, startnode, level = 0):
        self.counter += 1

        if startnode in self.checked:
            # Don't redo parts of tree already checked
            return

        self.checked.append(startnode)
        inputs = [startnode.input(x) for x in range(startnode.inputs()) if startnode.input(x) is not None]
        for i in inputs:
            edge = pydot.Edge(i.name(), startnode.name())
            self.graph.add_edge(edge)
            print "%d Recusing into %s" %  (level, i.name())
            self.recursive_grapher(i, level = level + 1)

    def go(self, startnode, fname):
        self.recursive_grapher(startnode)
        self.graph.write_dot(fname)

import re
def nodepos_from_dot(fname):
    for line in open(fname):
        line = line.strip()
        m = re.match("([\w\d_]+) \[pos=\"(\d+),(\d+)\",", line)
        if m is None:
            continue
        print m.groups()
        nodename, xpos, ypos = m.groups()
        n = nuke.toNode(nodename)
        n.setXpos(int(xpos))
        n.setYpos(-int(ypos))

a = NukeDoter()
fname = '/tmp/nuke.dot'
a.go(nuke.selectedNode(), fname)
nodepos_from_dot(fname)