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