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         "wiki_url": "",
  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()