1# ##### BEGIN GPL LICENSE BLOCK #####
2#
3#  This program is free software; you can redistribute it and/or
4#  modify it under the terms of the GNU General Public License
5#  as published by the Free Software Foundation; either version 2
6#  of the License, or (at your option) any later version.
7#
8#  This program is distributed in the hope that it will be useful,
9#  but WITHOUT ANY WARRANTY; without even the implied warranty of
10#  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
11#  GNU General Public License for more details.
12#
13#  You should have received a copy of the GNU General Public License
14#  along with this program; if not, write to the Free Software Foundation,
15#  Inc., 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301, USA.
16#
17# ##### END GPL LICENSE BLOCK #####
18
19bl_info = {
20    "name": "Node Arrange",
21    "author": "JuhaW",
22    "version": (0, 2, 2),
23    "blender": (2, 80, 4),
24    "location": "Node Editor > Properties > Trees",
25    "description": "Node Tree Arrangement Tools",
26    "warning": "",
27    "doc_url": "{BLENDER_MANUAL_URL}/addons/node/node_arrange.html",
28    "tracker_url": "https://github.com/JuhaW/NodeArrange/issues",
29    "category": "Node"
30}
31
32
33import sys
34import bpy
35from collections import OrderedDict
36from itertools import repeat
37import pprint
38import pdb
39from bpy.types import Operator, Panel
40from bpy.props import (
41    IntProperty,
42)
43from copy import copy
44
45
46#From Node Wrangler
47def get_nodes_linked(context):
48    tree = context.space_data.node_tree
49
50    # Get nodes from currently edited tree.
51    # If user is editing a group, space_data.node_tree is still the base level (outside group).
52    # context.active_node is in the group though, so if space_data.node_tree.nodes.active is not
53    # the same as context.active_node, the user is in a group.
54    # Check recursively until we find the real active node_tree:
55    if tree.nodes.active:
56        while tree.nodes.active != context.active_node:
57            tree = tree.nodes.active.node_tree
58
59    return tree.nodes, tree.links
60
61class NA_OT_AlignNodes(Operator):
62    '''Align the selected nodes/Tidy loose nodes'''
63    bl_idname = "node.na_align_nodes"
64    bl_label = "Align Nodes"
65    bl_options = {'REGISTER', 'UNDO'}
66    margin: IntProperty(name='Margin', default=50, description='The amount of space between nodes')
67
68    def execute(self, context):
69        nodes, links = get_nodes_linked(context)
70        margin = self.margin
71
72        selection = []
73        for node in nodes:
74            if node.select and node.type != 'FRAME':
75                selection.append(node)
76
77        # If no nodes are selected, align all nodes
78        active_loc = None
79        if not selection:
80            selection = nodes
81        elif nodes.active in selection:
82            active_loc = copy(nodes.active.location)  # make a copy, not a reference
83
84        # Check if nodes should be laid out horizontally or vertically
85        x_locs = [n.location.x + (n.dimensions.x / 2) for n in selection]  # use dimension to get center of node, not corner
86        y_locs = [n.location.y - (n.dimensions.y / 2) for n in selection]
87        x_range = max(x_locs) - min(x_locs)
88        y_range = max(y_locs) - min(y_locs)
89        mid_x = (max(x_locs) + min(x_locs)) / 2
90        mid_y = (max(y_locs) + min(y_locs)) / 2
91        horizontal = x_range > y_range
92
93        # Sort selection by location of node mid-point
94        if horizontal:
95            selection = sorted(selection, key=lambda n: n.location.x + (n.dimensions.x / 2))
96        else:
97            selection = sorted(selection, key=lambda n: n.location.y - (n.dimensions.y / 2), reverse=True)
98
99        # Alignment
100        current_pos = 0
101        for node in selection:
102            current_margin = margin
103            current_margin = current_margin * 0.5 if node.hide else current_margin  # use a smaller margin for hidden nodes
104
105            if horizontal:
106                node.location.x = current_pos
107                current_pos += current_margin + node.dimensions.x
108                node.location.y = mid_y + (node.dimensions.y / 2)
109            else:
110                node.location.y = current_pos
111                current_pos -= (current_margin * 0.3) + node.dimensions.y  # use half-margin for vertical alignment
112                node.location.x = mid_x - (node.dimensions.x / 2)
113
114        # If active node is selected, center nodes around it
115        if active_loc is not None:
116            active_loc_diff = active_loc - nodes.active.location
117            for node in selection:
118                node.location += active_loc_diff
119        else:  # Position nodes centered around where they used to be
120            locs = ([n.location.x + (n.dimensions.x / 2) for n in selection]) if horizontal else ([n.location.y - (n.dimensions.y / 2) for n in selection])
121            new_mid = (max(locs) + min(locs)) / 2
122            for node in selection:
123                if horizontal:
124                    node.location.x += (mid_x - new_mid)
125                else:
126                    node.location.y += (mid_y - new_mid)
127
128        return {'FINISHED'}
129
130class values():
131    average_y = 0
132    x_last = 0
133    margin_x = 100
134    mat_name = ""
135    margin_y = 20
136
137
138class NA_PT_NodePanel(Panel):
139    bl_label = "Node Arrange"
140    bl_space_type = "NODE_EDITOR"
141    bl_region_type = "UI"
142    bl_category = "Arrange"
143
144    def draw(self, context):
145        if context.active_node is not None:
146            layout = self.layout
147            row = layout.row()
148            col = layout.column
149            row.operator('node.button')
150
151            row = layout.row()
152            row.prop(bpy.context.scene, 'nodemargin_x', text="Margin x")
153            row = layout.row()
154            row.prop(bpy.context.scene, 'nodemargin_y', text="Margin y")
155            row = layout.row()
156            row.prop(context.scene, 'node_center', text="Center nodes")
157
158            row = layout.row()
159            row.operator('node.na_align_nodes', text="Align to Selected")
160
161            row = layout.row()
162            node = context.space_data.node_tree.nodes.active
163            if node and node.select:
164                row.prop(node, 'location', text = "Node X", index = 0)
165                row.prop(node, 'location', text = "Node Y", index = 1)
166                row = layout.row()
167                row.prop(node, 'width', text = "Node width")
168
169            row = layout.row()
170            row.operator('node.button_odd')
171
172class NA_OT_NodeButton(Operator):
173
174    '''Arrange Connected Nodes/Arrange All Nodes'''
175    bl_idname = 'node.button'
176    bl_label = 'Arrange All Nodes'
177
178    def execute(self, context):
179        nodemargin(self, context)
180        bpy.context.space_data.node_tree.nodes.update()
181        bpy.ops.node.view_all()
182
183        return {'FINISHED'}
184
185    # not sure this is doing what you expect.
186    # blender.org/api/blender_python_api_current/bpy.types.Operator.html#invoke
187    def invoke(self, context, value):
188        values.mat_name = bpy.context.space_data.node_tree
189        nodemargin(self, context)
190        return {'FINISHED'}
191
192
193class NA_OT_NodeButtonOdd(Operator):
194
195    'Show the nodes for this material'
196    bl_idname = 'node.button_odd'
197    bl_label = 'Select Unlinked'
198
199    def execute(self, context):
200        values.mat_name = bpy.context.space_data.node_tree
201        #mat = bpy.context.object.active_material
202        nodes_iterate(context.space_data.node_tree, False)
203        return {'FINISHED'}
204
205
206class NA_OT_NodeButtonCenter(Operator):
207
208    'Show the nodes for this material'
209    bl_idname = 'node.button_center'
210    bl_label = 'Center nodes (0,0)'
211
212    def execute(self, context):
213        values.mat_name = ""  # reset
214        mat = bpy.context.object.active_material
215        nodes_center(mat)
216        return {'FINISHED'}
217
218
219def nodemargin(self, context):
220
221    values.margin_x = context.scene.nodemargin_x
222    values.margin_y = context.scene.nodemargin_y
223
224    ntree = context.space_data.node_tree
225
226    #first arrange nodegroups
227    n_groups = []
228    for i in ntree.nodes:
229        if i.type == 'GROUP':
230            n_groups.append(i)
231
232    while n_groups:
233        j = n_groups.pop(0)
234        nodes_iterate(j.node_tree)
235        for i in j.node_tree.nodes:
236            if i.type == 'GROUP':
237                n_groups.append(i)
238
239    nodes_iterate(ntree)
240
241    # arrange nodes + this center nodes together
242    if context.scene.node_center:
243        nodes_center(ntree)
244
245
246class NA_OT_ArrangeNodesOp(bpy.types.Operator):
247    bl_idname = 'node.arrange_nodetree'
248    bl_label = 'Nodes Private Op'
249
250    mat_name : bpy.props.StringProperty()
251    margin_x : bpy.props.IntProperty(default=120)
252    margin_y : bpy.props.IntProperty(default=120)
253
254    def nodemargin2(self, context):
255        mat = None
256        mat_found = bpy.data.materials.get(self.mat_name)
257        if self.mat_name and mat_found:
258            mat = mat_found
259            #print(mat)
260
261        if not mat:
262            return
263        else:
264            values.mat_name = self.mat_name
265            scn = context.scene
266            scn.nodemargin_x = self.margin_x
267            scn.nodemargin_y = self.margin_y
268            nodes_iterate(mat)
269            if scn.node_center:
270                nodes_center(mat)
271
272    def execute(self, context):
273        self.nodemargin2(context)
274        return {'FINISHED'}
275
276
277def outputnode_search(ntree):    # return node/None
278
279    outputnodes = []
280    for node in ntree.nodes:
281        if not node.outputs:
282            for input in node.inputs:
283                if input.is_linked:
284                    outputnodes.append(node)
285                    break
286
287    if not outputnodes:
288        print("No output node found")
289        return None
290    return outputnodes
291
292
293###############################################################
294def nodes_iterate(ntree, arrange=True):
295
296    nodeoutput = outputnode_search(ntree)
297    if nodeoutput is None:
298        #print ("nodeoutput is None")
299        return None
300    a = []
301    a.append([])
302    for i in nodeoutput:
303        a[0].append(i)
304
305
306    level = 0
307
308    while a[level]:
309        a.append([])
310
311        for node in a[level]:
312            inputlist = [i for i in node.inputs if i.is_linked]
313
314            if inputlist:
315
316                for input in inputlist:
317                    for nlinks in input.links:
318                        node1 = nlinks.from_node
319                        a[level + 1].append(node1)
320
321            else:
322                pass
323
324        level += 1
325
326    del a[level]
327    level -= 1
328
329    #remove duplicate nodes at the same level, first wins
330    for x, nodes in enumerate(a):
331        a[x] = list(OrderedDict(zip(a[x], repeat(None))))
332
333    #remove duplicate nodes in all levels, last wins
334    top = level
335    for row1 in range(top, 1, -1):
336        for col1 in a[row1]:
337            for row2 in range(row1-1, 0, -1):
338                for col2 in a[row2]:
339                    if col1 == col2:
340                        a[row2].remove(col2)
341                        break
342
343    """
344    for x, i in enumerate(a):
345        print (x)
346        for j in i:
347            print (j)
348        #print()
349    """
350    """
351    #add node frames to nodelist
352    frames = []
353    print ("Frames:")
354    print ("level:", level)
355    print ("a:",a)
356    for row in range(level, 0, -1):
357
358        for i, node in enumerate(a[row]):
359            if node.parent:
360                print ("Frame found:", node.parent, node)
361                #if frame already added to the list ?
362                frame = node.parent
363                #remove node
364                del a[row][i]
365                if frame not in frames:
366                    frames.append(frame)
367                    #add frame to the same place than node was
368                    a[row].insert(i, frame)
369
370    pprint.pprint(a)
371    """
372    #return None
373    ########################################
374
375
376
377    if not arrange:
378        nodelist = [j for i in a for j in i]
379        nodes_odd(ntree, nodelist=nodelist)
380        return None
381
382    ########################################
383
384    levelmax = level + 1
385    level = 0
386    values.x_last = 0
387
388    while level < levelmax:
389
390        values.average_y = 0
391        nodes = [x for x in a[level]]
392        #print ("level, nodes:", level, nodes)
393        nodes_arrange(nodes, level)
394
395        level = level + 1
396
397    return None
398
399
400###############################################################
401def nodes_odd(ntree, nodelist):
402
403    nodes = ntree.nodes
404    for i in nodes:
405        i.select = False
406
407    a = [x for x in nodes if x not in nodelist]
408    # print ("odd nodes:",a)
409    for i in a:
410        i.select = True
411
412
413def nodes_arrange(nodelist, level):
414
415    parents = []
416    for node in nodelist:
417        parents.append(node.parent)
418        node.parent = None
419        bpy.context.space_data.node_tree.nodes.update()
420
421
422    #print ("nodes arrange def")
423    # node x positions
424
425    widthmax = max([x.dimensions.x for x in nodelist])
426    xpos = values.x_last - (widthmax + values.margin_x) if level != 0 else 0
427    #print ("nodelist, xpos", nodelist,xpos)
428    values.x_last = xpos
429
430    # node y positions
431    x = 0
432    y = 0
433
434    for node in nodelist:
435
436        if node.hide:
437            hidey = (node.dimensions.y / 2) - 8
438            y = y - hidey
439        else:
440            hidey = 0
441
442        node.location.y = y
443        y = y - values.margin_y - node.dimensions.y + hidey
444
445        node.location.x = xpos #if node.type != "FRAME" else xpos + 1200
446
447    y = y + values.margin_y
448
449    center = (0 + y) / 2
450    values.average_y = center - values.average_y
451
452    #for node in nodelist:
453
454        #node.location.y -= values.average_y
455
456    for i, node in enumerate(nodelist):
457        node.parent =  parents[i]
458
459def nodetree_get(mat):
460
461    return mat.node_tree.nodes
462
463
464def nodes_center(ntree):
465
466    bboxminx = []
467    bboxmaxx = []
468    bboxmaxy = []
469    bboxminy = []
470
471    for node in ntree.nodes:
472        if not node.parent:
473            bboxminx.append(node.location.x)
474            bboxmaxx.append(node.location.x + node.dimensions.x)
475            bboxmaxy.append(node.location.y)
476            bboxminy.append(node.location.y - node.dimensions.y)
477
478    # print ("bboxminy:",bboxminy)
479    bboxminx = min(bboxminx)
480    bboxmaxx = max(bboxmaxx)
481    bboxminy = min(bboxminy)
482    bboxmaxy = max(bboxmaxy)
483    center_x = (bboxminx + bboxmaxx) / 2
484    center_y = (bboxminy + bboxmaxy) / 2
485    '''
486    print ("minx:",bboxminx)
487    print ("maxx:",bboxmaxx)
488    print ("miny:",bboxminy)
489    print ("maxy:",bboxmaxy)
490
491    print ("bboxes:", bboxminx, bboxmaxx, bboxmaxy, bboxminy)
492    print ("center x:",center_x)
493    print ("center y:",center_y)
494    '''
495
496    x = 0
497    y = 0
498
499    for node in ntree.nodes:
500
501        if not node.parent:
502            node.location.x -= center_x
503            node.location.y += -center_y
504
505classes = [
506    NA_PT_NodePanel,
507    NA_OT_NodeButton,
508    NA_OT_NodeButtonOdd,
509    NA_OT_NodeButtonCenter,
510    NA_OT_ArrangeNodesOp,
511    NA_OT_AlignNodes
512]
513
514def register():
515    for c in classes:
516        bpy.utils.register_class(c)
517
518    bpy.types.Scene.nodemargin_x = bpy.props.IntProperty(default=100, update=nodemargin)
519    bpy.types.Scene.nodemargin_y = bpy.props.IntProperty(default=20, update=nodemargin)
520    bpy.types.Scene.node_center = bpy.props.BoolProperty(default=True, update=nodemargin)
521
522
523
524def unregister():
525    for c in classes:
526        bpy.utils.unregister_class(c)
527
528    del bpy.types.Scene.nodemargin_x
529    del bpy.types.Scene.nodemargin_y
530    del bpy.types.Scene.node_center
531
532if __name__ == "__main__":
533    register()
534