node_arrange.py

   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
  19 bl_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
  33 import sys
  34 import bpy
  35 from collections import OrderedDict
  36 from itertools import repeat
  37 import pprint
  38 import pdb
  39 from bpy.types import Operator, Panel
  40 from bpy.props import (
  41     IntProperty,
  42 )
  43 from copy import copy
  44
  45
  46 #From Node Wrangler
  47 def 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
  61 class 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
 130 class values():
 131     average_y = 0
 132     x_last = 0
 133     margin_x = 100
 134     mat_name = ""
 135     margin_y = 20
 136
 137
 138 class 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
 172 class 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
 193 class 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
 206 class 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
 219 def 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
 246 class 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
 277 def 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 ###############################################################
 294 def 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 ###############################################################
 401 def 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
 413 def 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
 459 def nodetree_get(mat):
 460
 461     return mat.node_tree.nodes
 462
 463
 464 def 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
 505 classes = [
 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
 514 def 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
 524 def 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
 532 if __name__ == "__main__":
 533     register()