node_arrange.py

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