Node Arrange: 2.81, Thanks JuhaW: T66410

node_arrange.py

+# ##### BEGIN GPL LICENSE BLOCK #####
+#
+#  This program is free software; you can redistribute it and/or
+#  modify it under the terms of the GNU General Public License
+#  as published by the Free Software Foundation; either version 2
+#  of the License, or (at your option) any later version.
+#
+#  This program is distributed in the hope that it will be useful,
+#  but WITHOUT ANY WARRANTY; without even the implied warranty of
+#  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+#  GNU General Public License for more details.
+#
+#  You should have received a copy of the GNU General Public License
+#  along with this program; if not, write to the Free Software Foundation,
+#  Inc., 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301, USA.
+#
+# ##### END GPL LICENSE BLOCK #####
+
+bl_info = {
+	"name": "Node Arrange",
+	"author": "JuhaW",
+	"version": (0, 2, 1),
+	"blender": (2, 80, 4),
+	"location": "Node Editor > Properties > Trees",
+	"description": "Node Tree Arrangement Tools",
+	"warning": "",
+	"wiki_url": "",
+	"tracker_url": "https://github.com/JuhaW/NodeArrange/issues",
+	"category": "Node"
+}
+
+
+import sys
+import bpy
+from collections import OrderedDict
+from itertools import repeat
+import pprint
+import pdb
+from bpy.types import Operator, Panel
+from bpy.props import (
+    IntProperty,
+)
+from copy import copy
+
+
+#From Node Wrangler
+def get_nodes_linked(context):
+    tree = context.space_data.node_tree
+
+    # Get nodes from currently edited tree.
+    # If user is editing a group, space_data.node_tree is still the base level (outside group).
+    # context.active_node is in the group though, so if space_data.node_tree.nodes.active is not
+    # the same as context.active_node, the user is in a group.
+    # Check recursively until we find the real active node_tree:
+    if tree.nodes.active:
+        while tree.nodes.active != context.active_node:
+            tree = tree.nodes.active.node_tree
+
+    return tree.nodes, tree.links
+
+class NA_OT_AlignNodes(Operator):
+    '''Align the selected nodes/Tidy loose nodes'''
+    bl_idname = "node.na_align_nodes"
+    bl_label = "Align Nodes"
+    bl_options = {'REGISTER', 'UNDO'}
+    margin: IntProperty(name='Margin', default=50, description='The amount of space between nodes')
+
+    def execute(self, context):
+        nodes, links = get_nodes_linked(context)
+        margin = self.margin
+
+        selection = []
+        for node in nodes:
+            if node.select and node.type != 'FRAME':
+                selection.append(node)
+
+        # If no nodes are selected, align all nodes
+        active_loc = None
+        if not selection:
+            selection = nodes
+        elif nodes.active in selection:
+            active_loc = copy(nodes.active.location)  # make a copy, not a reference
+
+        # Check if nodes should be laid out horizontally or vertically
+        x_locs = [n.location.x + (n.dimensions.x / 2) for n in selection]  # use dimension to get center of node, not corner
+        y_locs = [n.location.y - (n.dimensions.y / 2) for n in selection]
+        x_range = max(x_locs) - min(x_locs)
+        y_range = max(y_locs) - min(y_locs)
+        mid_x = (max(x_locs) + min(x_locs)) / 2
+        mid_y = (max(y_locs) + min(y_locs)) / 2
+        horizontal = x_range > y_range
+
+        # Sort selection by location of node mid-point
+        if horizontal:
+            selection = sorted(selection, key=lambda n: n.location.x + (n.dimensions.x / 2))
+        else:
+            selection = sorted(selection, key=lambda n: n.location.y - (n.dimensions.y / 2), reverse=True)
+
+        # Alignment
+        current_pos = 0
+        for node in selection:
+            current_margin = margin
+            current_margin = current_margin * 0.5 if node.hide else current_margin  # use a smaller margin for hidden nodes
+
+            if horizontal:
+                node.location.x = current_pos
+                current_pos += current_margin + node.dimensions.x
+                node.location.y = mid_y + (node.dimensions.y / 2)
+            else:
+                node.location.y = current_pos
+                current_pos -= (current_margin * 0.3) + node.dimensions.y  # use half-margin for vertical alignment
+                node.location.x = mid_x - (node.dimensions.x / 2)
+
+        # If active node is selected, center nodes around it
+        if active_loc is not None:
+            active_loc_diff = active_loc - nodes.active.location
+            for node in selection:
+                node.location += active_loc_diff
+        else:  # Position nodes centered around where they used to be
+            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])
+            new_mid = (max(locs) + min(locs)) / 2
+            for node in selection:
+                if horizontal:
+                    node.location.x += (mid_x - new_mid)
+                else:
+                    node.location.y += (mid_y - new_mid)
+
+        return {'FINISHED'}
+
+class values():
+	average_y = 0
+	x_last = 0
+	margin_x = 100
+	mat_name = ""
+	margin_y = 20
+
+
+class NA_PT_NodePanel(Panel):
+	bl_label = "Node Arrange"
+	bl_space_type = "NODE_EDITOR"
+	bl_region_type = "UI"
+	bl_category = "Arrange"
+
+	def draw(self, context):
+		if context.active_node is not None:
+			layout = self.layout
+			row = layout.row()
+			col = layout.column
+			row.operator('node.button')
+
+			row = layout.row()
+			row.prop(bpy.context.scene, 'nodemargin_x', text="Margin x")
+			row = layout.row()
+			row.prop(bpy.context.scene, 'nodemargin_y', text="Margin y")
+			row = layout.row()
+			row.prop(context.scene, 'node_center', text="Center nodes")
+
+			row = layout.row()
+			row.operator('node.na_align_nodes', text="Align to Selected")
+
+			row = layout.row()
+			node = context.space_data.node_tree.nodes.active
+			if node and node.select:
+				row.prop(node, 'location', text = "Node X", index = 0)
+				row.prop(node, 'location', text = "Node Y", index = 1)
+				row = layout.row()
+				row.prop(node, 'width', text = "Node width")
+
+			row = layout.row()
+			row.operator('node.button_odd')
+
+class NA_OT_NodeButton(Operator):
+
+	'''Arrange Connected Nodes/Arrange All Nodes'''
+	bl_idname = 'node.button'
+	bl_label = 'Arrange All Nodes'
+
+	def execute(self, context):
+		nodemargin(self, context)
+		bpy.context.space_data.node_tree.nodes.update()
+		bpy.ops.node.view_all()
+
+		return {'FINISHED'}
+
+	# not sure this is doing what you expect.
+	# blender.org/api/blender_python_api_current/bpy.types.Operator.html#invoke
+	def invoke(self, context, value):
+		values.mat_name = bpy.context.space_data.node_tree
+		nodemargin(self, context)
+		return {'FINISHED'}
+
+
+class NA_OT_NodeButtonOdd(Operator):
+
+	'Show the nodes for this material'
+	bl_idname = 'node.button_odd'
+	bl_label = 'Select Unlinked'
+
+	def execute(self, context):
+		values.mat_name = bpy.context.space_data.node_tree
+		#mat = bpy.context.object.active_material
+		nodes_iterate(context.space_data.node_tree, False)
+		return {'FINISHED'}
+
+
+class NA_OT_NodeButtonCenter(Operator):
+
+	'Show the nodes for this material'
+	bl_idname = 'node.button_center'
+	bl_label = 'Center nodes (0,0)'
+
+	def execute(self, context):
+		values.mat_name = ""  # reset
+		mat = bpy.context.object.active_material
+		nodes_center(mat)
+		return {'FINISHED'}
+
+
+def nodemargin(self, context):
+
+	values.margin_x = context.scene.nodemargin_x
+	values.margin_y = context.scene.nodemargin_y
+
+	ntree = context.space_data.node_tree
+
+	nodes_iterate(ntree)
+
+	# arrange nodes + this center nodes together
+	if context.scene.node_center:
+		nodes_center(ntree)
+
+
+class NA_OT_ArrangeNodesOp(bpy.types.Operator):
+	bl_idname = 'node.arrange_nodetree'
+	bl_label = 'Nodes Private Op'
+
+	mat_name : bpy.props.StringProperty()
+	margin_x : bpy.props.IntProperty(default=120)
+	margin_y : bpy.props.IntProperty(default=120)
+
+	def nodemargin2(self, context):
+		mat = None
+		mat_found = bpy.data.materials.get(self.mat_name)
+		if self.mat_name and mat_found:
+			mat = mat_found
+			#print(mat)
+
+		if not mat:
+			return
+		else:
+			values.mat_name = self.mat_name
+			scn = context.scene
+			scn.nodemargin_x = self.margin_x
+			scn.nodemargin_y = self.margin_y
+			nodes_iterate(mat)
+			if scn.node_center:
+				nodes_center(mat)
+
+	def execute(self, context):
+		self.nodemargin2(context)
+		return {'FINISHED'}
+
+
+def outputnode_search(ntree):	 # return node/None
+
+	outputnodes = []
+	for node in bpy.context.space_data.node_tree.nodes:
+		if not node.outputs:
+			for input in node.inputs:
+				if input.is_linked:
+					outputnodes.append(node)
+					break
+
+	if not outputnodes:
+		print("No output node found")
+		return None
+	return outputnodes
+
+
+###############################################################
+def nodes_iterate(ntree, arrange=True):
+
+	nodeoutput = outputnode_search(ntree)
+	if nodeoutput is None:
+		#print ("nodeoutput is None")
+		return None
+	a = []
+	a.append([])
+	for i in nodeoutput:
+		a[0].append(i)
+
+
+	level = 0
+
+	while a[level]:
+		a.append([])
+
+		for node in a[level]:
+			inputlist = [i for i in node.inputs if i.is_linked]
+
+			if inputlist:
+
+				for input in inputlist:
+					for nlinks in input.links:
+						node1 = nlinks.from_node
+						a[level + 1].append(node1)
+
+			else:
+				pass
+
+		level += 1
+
+	del a[level]
+	level -= 1
+
+	#remove duplicate nodes at the same level, first wins
+	for x, nodes in enumerate(a):
+		a[x] = list(OrderedDict(zip(a[x], repeat(None))))
+
+	#remove duplicate nodes in all levels, last wins
+	top = level
+	for row1 in range(top, 1, -1):
+		for col1 in a[row1]:
+			for row2 in range(row1-1, 0, -1):
+				for col2 in a[row2]:
+					if col1 == col2:
+						a[row2].remove(col2)
+						break
+
+	"""
+	for x, i in enumerate(a):
+		print (x)
+		for j in i:
+			print (j)
+		#print()
+	"""
+	"""
+	#add node frames to nodelist
+	frames = []
+	print ("Frames:")
+	print ("level:", level)
+	print ("a:",a)
+	for row in range(level, 0, -1):
+
+		for i, node in enumerate(a[row]):
+			if node.parent:
+				print ("Frame found:", node.parent, node)
+				#if frame already added to the list ?
+				frame = node.parent
+				#remove node
+				del a[row][i]
+				if frame not in frames:
+					frames.append(frame)
+					#add frame to the same place than node was
+					a[row].insert(i, frame)
+
+	pprint.pprint(a)
+	"""
+	#return None
+	########################################
+
+
+
+	if not arrange:
+		nodelist = [j for i in a for j in i]
+		nodes_odd(ntree, nodelist=nodelist)
+		return None
+
+	########################################
+
+	levelmax = level + 1
+	level = 0
+	values.x_last = 0
+
+	while level < levelmax:
+
+		values.average_y = 0
+		nodes = [x for x in a[level]]
+		#print ("level, nodes:", level, nodes)
+		nodes_arrange(nodes, level)
+
+		level = level + 1
+
+	return None
+
+
+###############################################################
+def nodes_odd(ntree, nodelist):
+
+	nodes = ntree.nodes
+	for i in nodes:
+		i.select = False
+
+	a = [x for x in nodes if x not in nodelist]
+	# print ("odd nodes:",a)
+	for i in a:
+		i.select = True
+
+
+def nodes_arrange(nodelist, level):
+
+	parents = []
+	for node in nodelist:
+		parents.append(node.parent)
+		node.parent = None
+		bpy.context.space_data.node_tree.nodes.update()
+
+
+	#print ("nodes arrange def")
+	# node x positions
+
+	widthmax = max([x.dimensions.x for x in nodelist])
+	xpos = values.x_last - (widthmax + values.margin_x) if level != 0 else 0
+	#print ("nodelist, xpos", nodelist,xpos)
+	values.x_last = xpos
+
+	# node y positions
+	x = 0
+	y = 0
+
+	for node in nodelist:
+
+		if node.hide:
+			hidey = (node.dimensions.y / 2) - 8
+			y = y - hidey
+		else:
+			hidey = 0
+
+		node.location.y = y
+		y = y - values.margin_y - node.dimensions.y + hidey
+
+		node.location.x = xpos #if node.type != "FRAME" else xpos + 1200
+
+	y = y + values.margin_y
+
+	center = (0 + y) / 2
+	values.average_y = center - values.average_y
+
+	#for node in nodelist:
+
+		#node.location.y -= values.average_y
+
+	for i, node in enumerate(nodelist):
+		node.parent =  parents[i]
+
+def nodetree_get(mat):
+
+	return mat.node_tree.nodes
+
+
+def nodes_center(ntree):
+
+	bboxminx = []
+	bboxmaxx = []
+	bboxmaxy = []
+	bboxminy = []
+
+	for node in ntree.nodes:
+		if not node.parent:
+			bboxminx.append(node.location.x)
+			bboxmaxx.append(node.location.x + node.dimensions.x)
+			bboxmaxy.append(node.location.y)
+			bboxminy.append(node.location.y - node.dimensions.y)
+
+	# print ("bboxminy:",bboxminy)
+	bboxminx = min(bboxminx)
+	bboxmaxx = max(bboxmaxx)
+	bboxminy = min(bboxminy)
+	bboxmaxy = max(bboxmaxy)
+	center_x = (bboxminx + bboxmaxx) / 2
+	center_y = (bboxminy + bboxmaxy) / 2
+	'''
+	print ("minx:",bboxminx)
+	print ("maxx:",bboxmaxx)
+	print ("miny:",bboxminy)
+	print ("maxy:",bboxmaxy)
+
+	print ("bboxes:", bboxminx, bboxmaxx, bboxmaxy, bboxminy)
+	print ("center x:",center_x)
+	print ("center y:",center_y)
+	'''
+
+	x = 0
+	y = 0
+
+	for node in ntree.nodes:
+
+		if not node.parent:
+			node.location.x -= center_x
+			node.location.y += -center_y
+
+classes = [
+	NA_PT_NodePanel,
+	NA_OT_NodeButton,
+	NA_OT_NodeButtonOdd,
+	NA_OT_NodeButtonCenter,
+	NA_OT_ArrangeNodesOp,
+    NA_OT_AlignNodes
+]
+
+def register():
+	for c in classes:
+		bpy.utils.register_class(c)
+
+	bpy.types.Scene.nodemargin_x = bpy.props.IntProperty(default=100, update=nodemargin)
+	bpy.types.Scene.nodemargin_y = bpy.props.IntProperty(default=20, update=nodemargin)
+	bpy.types.Scene.node_center = bpy.props.BoolProperty(default=True, update=nodemargin)
+
+
+
+def unregister():
+	for c in classes:
+		bpy.utils.unregister_class(c)
+
+	del bpy.types.Scene.nodemargin_x
+	del bpy.types.Scene.nodemargin_y
+	del bpy.types.Scene.node_center
+
+if __name__ == "__main__":
+	register()