Elbows

1# -*- coding: utf-8 -*-
2"""
3Nuke Elbows
4
5This script create right angle 'elbows' with the dot node.
6
7Installation: 
8Add the following to menu.py
9
10$ import elbows
11$ nuke.menu("Nuke").addCommand('Scripts/Elbows', 'elbows.main()', "shift+alt+e")
12
13"""
14
15
16import nuke
17
18
19
20def nodeHasMask(node):
21    return node.knobs().has_key('maskChannelMask')
22
23
24def getMinInputs(node):
25    min_inputs = node.minInputs()
26    return (min_inputs - 1) if nodeHasMask(node) and min_inputs > 0 else min_inputs
27
28
29def getPos(node):
30    return node.xpos(), node.ypos()
31
32
33def getSize(node):
34    return node.screenWidth() / 2, node.screenHeight() / 2
35
36
37def checkUnder(src, dst):
38    # True if source is within 1/2 the max of both src and dst widths on x axis
39    m = getPos(src)[0] + getSize(src)[0]
40    d = getPos(dst)[0] + getSize(dst)[0]
41    threshold = max(getSize(src)[0], getSize(dst)[0])  # threshold is max of half widths
42    return abs(m - d) < threshold
43
44
45def align(dot, dst, mer):
46    # Aligns dots in a parallel way
47    size_dot_x, size_dot_y = getSize(dot)
48    x = getPos(dst)[0] + getSize(dst)[0] - size_dot_x
49    y = getPos(mer)[1] + getSize(mer)[1] - size_dot_y
50    dot.setXYpos(x, y)
51    return x, y
52
53
54def align_overhead(dot, dep, node):  # src = dep
55    # Create and align dot overhead
56    size_dot_x, size_dot_y = getSize(dot)
57    x = getPos(node)[0] + getSize(node)[0] - size_dot_x
58    y = getPos(dep)[1] + getSize(dep)[1] - size_dot_y
59    dot.setXYpos(x, y)
60    return x, y
61
62
63def elbows(node):
64    if len(node.dependencies(nuke.INPUTS)) == 1:
65
66        # The node has a single input, so do an overhead elbow
67
68        dep = node.input(0)
69        if dep is None:
70            return  # It's not connected to anything, exit
71
72        if not checkUnder(node, dep):  # Don't do anything if we're already under it
73            if dep.Class() == 'Dot':  # If its a dot, don't make another one
74                dot = nuke.nodes.Dot()
75                dot.setInput(0, dep)
76                align_overhead(dot, dep, node)
77                node.setInput(0, dot)
78                return [dot]
79
80        else:
81            node.setSelected(False)
82            dep.setSelected(True)
83            dot = nuke.createNode('Dot', inpanel=False)
84            dot.setYpos(dot.ypos() + node.screenHeight() * 2)
85            dot2 = nuke.nodes.Dot()
86            align_overhead(dot2, dot, node)
87            dot2.setInput(0, dot)
88            node.setInput(0, dot2)
89            return [dot, dot2]
90
91    elif len(node.dependencies(nuke.INPUTS)) > 1:
92
93        # The node has multiple inputs, create parallel elbows
94
95        if node.Class() in ('ContactSheet', 'Switch', 'Merge2',
96                            'Merge', 'Blend', 'Dissolve', 'Keymix',
97                            'Copy', 'AddMix', 'CopyBBox', 'ChannelMerge'):
98            dots = []
99            try:
100                child = node.dependent(nuke.INPUTS)[0]  # check if there is a connected node
101            except:
102                child = None
103            for input in range(node.inputs()):
104                if node.Class() not in ('ContactSheet', 'Switch'):
105                    if input in [2]:  # input 2 is the mask, skip it
106                        continue
107                dep = node.input(input)
108                if dep is not None:
109                    if not checkUnder(node, dep):
110                        dep.setSelected(True)
111                        dot = nuke.createNode('Dot', inpanel=False)
112                        dots.append(dot)
113                        dep.setSelected(False)
114                        align(dot, dep, node)
115                        dot.setSelected(False)
116                if child is not None:
117                    child.setInput(0, node)
118            return dots
119
120
121def main():
122    nodes = nuke.selectedNodes()
123    if len(nodes) == 1:
124        # A single node selection assumes an elbow joint on non mask inputs.
125        elbows(nodes[0])
126
127    elif len(nodes) > 1:
128        # If multiple nodes are selected, it assumes we want to add them together.
129        # empty lists are false, checking to see if these nodes are not already connected
130        if all(not n.dependent(nuke.INPUTS) for n in nodes):
131            m = nuke.createNode('Merge2')
132            last = nodes[-1]  # get the node that the merge should go under
133            m['operation'].setValue('plus')
134            m['bbox'].setValue('union')
135            m.setYpos(max(n.ypos() for n in nodes) + m.screenHeight() * 3)  # set Ypos under the lowest of selection
136            m.setXpos(getPos(last)[0] + getSize(last)[0] - getSize(m)[0])  # set Xpos to the last selected
137
138            for i, node in enumerate(nodes):
139                if i >= 2:
140                    i += 1
141                m.setInput(i, node)
142            dots = elbows(m)
143            for d in dots:
144                d.setSelected(True)
145            m.setSelected(True)