node_arrange.py

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