node_wrangler.py
Public
Actions

Authored By

	Greg Zaal (gregzaal)
	Nov 13 2013, 5:22 PM

Size

21 KB

Subscribers

None

node_wrangler.py
View Options

	# ##### 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 Wrangler",
	"author": "Greg Zaal",
	"version": (1, 0),
	"blender": (2, 67, 1),
	"location": "Node Editor > Z key/Q key or Properties Region",
	"description": "A set of tools that help clean up a node tree and improve viewing usability.",
	"warning": "",
	"wiki_url": "",
	"tracker_url": "",
	"category": "Node"}

	import bpy

	def actualRes(context, type="x"):
	rend=context.scene.render
	rend_percent=rend.resolution_percentage*0.01
	x=(str(rend.resolution_x*rend_percent).split('.'))[0]
	y=(str(rend.resolution_y*rend_percent).split('.'))[0]
	returned=y
	if type=="x":
	returned=x
	return (returned)

	def get_nodes_links(context): # Taken from Node Efficiency Tools by Bartek Skorupa
	space = context.space_data
	tree = space.node_tree
	nodes = tree.nodes
	links = tree.links
	active = nodes.active
	context_active = context.active_node
	# check if we are working on regular node tree or node group is currently edited.
	# if group is edited - active node of space_tree is the group
	# if context.active_node != space active node - it means that the group is being edited.
	# in such case we set "nodes" to be nodes of this group, "links" to be links of this group
	# if context.active_node == space.active_node it means that we are not currently editing group
	is_main_tree = True
	if active:
	is_main_tree = context_active == active
	if not is_main_tree: # if group is currently edited
	tree = active.node_tree
	nodes = tree.nodes
	links = tree.links
	if context.scene.SelectionOnly:
	newnodes=[]
	for node in nodes:
	if node.select==True:
	newnodes.append(node)
	nodes=newnodes
	nodes_sorted=sorted(nodes, key=lambda x: x.name) # Sort the nodes list to achieve consistent
	links_sorted=sorted(links, key=lambda x: x.from_node.name) # results (order was changed based on selection).
	return nodes_sorted, links_sorted

	def isStartNode(node):
	bool=True
	if len(node.inputs):
	for input in node.inputs:
	if input.links != ():
	bool=False
	return bool

	def isEndNode(node):
	bool=True
	if len(node.outputs):
	for output in node.outputs:
	if output.links != ():
	bool=False
	return bool

	def between(b1, a, b2):
	# b1 MUST be smaller than b2!
	bool=False
	if a >= b1 and a<=b2:
	bool=True
	return bool

	def overlaps (node1, node2):
	dim1x=node1.dimensions.x
	dim1y=node1.dimensions.y
	dim2x=node2.dimensions.x
	dim2y=node2.dimensions.y
	boolx=False
	booly=False
	boolboth=False

	# check for x overlap
	if between(node2.location.x, node1.location.x, (node2.location.x+dim2x)) or between(node2.location.x, (node1.location.x+dim1x), (node2.location.x+dim2x)): #if either edges are inside the second node
	boolx=True
	if between(node1.location.x, node2.location.x, node1.location.x+dim1x) and between(node1.location.x, (node2.location.x+dim2x), node1.location.x+dim1x): #if each edge is on either side of the second node
	boolx=True

	# check for y overlap
	if between((node2.location.y-dim2y), node1.location.y, node2.location.y) or between((node2.location.y-dim2y), (node1.location.y-dim1y), node2.location.y):
	booly=True
	if between((node1.location.y-dim1y), node2.location.y, node1.location.y) and between((node1.location.y-dim1y), (node2.location.y-dim2y), node1.location.y):
	booly=True

	if boolx==True and booly==True:
	boolboth=True
	return boolboth

	def treeMidPt(nodes):
	minx=(sorted(nodes, key=lambda k: k.location.x))[0].location.x
	miny=(sorted(nodes, key=lambda k: k.location.y))[0].location.y
	maxx=(sorted(nodes, key=lambda k: k.location.x, reverse=True))[0].location.x
	maxy=(sorted(nodes, key=lambda k: k.location.y, reverse=True))[0].location.y

	midx=minx+((maxx-minx)/2)
	midy=miny+((maxy-miny)/2)

	return midx, midy

	class ArrangeNodes(bpy.types.Operator):
	'Automatically layout the node tree in a linear and non-overlapping fasion.'
	bl_idname='nodes.layout'
	bl_label='Arrange Nodes'

	def execute(self,context):
	print (" Arranging nodes...")
	nodes, links = get_nodes_links(context)
	margin=context.scene.Spacing

	oldmidx, oldmidy = treeMidPt(nodes)

	if context.scene.DelReroutes:
	# Store selection
	selection=[]
	for node in nodes:
	if node.select==True and node.type != "REROUTE":
	selection.append(node.name)
	#Delete Reroutes
	for node in nodes:
	node.select=False # deselect all nodes
	for node in nodes:
	if node.type == 'REROUTE':
	node.select=True
	bpy.ops.node.delete_reconnect()
	# Restore selection
	#nodes, links = get_nodes_links(context)
	for node in nodes:
	if node.name in selection:
	node.select=True
	else:
	# Store selection anyway
	selection=[]
	for node in nodes:
	if node.select==True:
	selection.append(node.name)

	if context.scene.FrameHandling=="delete":
	# Store selection
	selection=[]
	for node in nodes:
	if node.select==True and node.type != "FRAME":
	selection.append(node.name)
	#Delete Frames
	for node in nodes:
	node.select=False # deselect all nodes
	for node in nodes:
	if node.type == 'FRAME':
	node.select=True
	bpy.ops.node.delete()
	# Restore selection
	#nodes, links = get_nodes_links(context)
	for node in nodes:
	if node.name in selection:
	node.select=True

	layout_iterations = len(nodes)
	for it in range(0, layout_iterations):
	print (" Layout Iteration: "+str(it))
	for node in nodes:
	isframe=False
	if node.type=="FRAME" and context.scene.FrameHandling=='ignore':
	isframe=True
	if not isframe:
	if isStartNode(node) and context.scene.StartAlign: #line up start nodes
	node.location.x=node.dimensions.x/-2
	node.location.y=node.dimensions.y/2
	for link in links:
	if link.from_node == node and link.to_node in nodes:
	link.to_node.location.x=node.location.x+node.dimensions.x+margin
	link.to_node.location.y=node.location.y-(node.dimensions.y/2)+(link.to_node.dimensions.y/2)
	else:
	node.location.x=0
	node.location.y=0

	backward_check_iterations=len(nodes)
	for it in range(0, backward_check_iterations):
	for link in links:
	if link.from_node.location.x+link.from_node.dimensions.x >= link.to_node.location.x and link.to_node in nodes:
	link.to_node.location.x=link.from_node.location.x+link.from_node.dimensions.x+margin

	#line up end nodes
	if context.scene.EndAlign:
	for node in nodes:
	max_loc_x=(sorted(nodes, key=lambda x: x.location.x, reverse=True))[0].location.x
	if isEndNode(node) and not isStartNode(node):
	node.location.x=max_loc_x

	overlap_iterations = len(nodes)
	for it in range(0, overlap_iterations):
	for node in nodes:
	isframe=False
	if node.type=="FRAME" and context.scene.FrameHandling=='ignore':
	isframe=True
	if not isframe:
	for nodecheck in nodes:
	isframe=False
	if nodecheck.type=="FRAME" and context.scene.FrameHandling=='ignore':
	isframe=True
	if not isframe:
	if (node != nodecheck): #dont look for overlaps with self
	if overlaps(node, nodecheck):
	print (node.name + " overlaps "+nodecheck.name)
	node.location.y=nodecheck.location.y-nodecheck.dimensions.y-0.5*margin

	newmidx, newmidy = treeMidPt(nodes)
	middiffx=newmidx-oldmidx
	middiffy=newmidy-oldmidy

	# put nodes back to the center of the old center
	for node in nodes:
	node.location.x=node.location.x-middiffx
	node.location.y=node.location.y-middiffy

	print ()
	print ()
	return {'FINISHED'}

	class ArrangeSelectedNodes(bpy.types.Operator):
	'Automatically layout the selected nodes in a linear and non-overlapping fasion.'
	bl_idname='nodes.layout_sel'
	bl_label='Arrange Selected Nodes'

	def execute(self,context):
	context.scene.SelectionOnly=True
	bpy.ops.nodes.layout()
	context.scene.SelectionOnly=False
	return {'FINISHED'}

	class DeleteUnusedNodes(bpy.types.Operator):
	'Delete all nodes whose output is not used'
	bl_idname='nodes.del_unused'
	bl_label='Delete Unused Nodes'

	def execute(self,context):
	nodes, links = get_nodes_links(context)
	end_types=['OUTPUT_MATERIAL', 'OUTPUT', 'VIEWER', 'COMPOSITE', 'SPLITVIEWER', 'OUTPUT_FILE', 'LEVELS', 'OUTPUT_LAMP', 'OUTPUT_WORLD', 'GROUP', 'GROUP_INPUT', 'GROUP_OUTPUT']

	# Store selection
	selection=[]
	for node in nodes:
	if node.select==True:
	selection.append(node.name)

	deleted_nodes=[]
	del_unused_iterations = 5#len(nodes)
	for it in range(0, del_unused_iterations):
	for node in nodes:
	node.select=False # deselect all nodes
	for node in nodes:
	if isEndNode(node) and not node.type in end_types:
	node.select=True
	deleted_nodes.append(node.name)
	bpy.ops.node.delete()

	deleted_nodes=list(set(deleted_nodes)) # get unique list of deleted nodes (iterations would count the same node more than once)
	for n in deleted_nodes:
	self.report({'INFO'}, "Node "+n+" deleted")
	self.report({'INFO'}, "Deleted "+str(len(deleted_nodes))+" nodes")

	# Restore selection
	nodes, links = get_nodes_links(context)
	for node in nodes:
	if node.name in selection:
	node.select=True
	return {'FINISHED'}

	class GetImage(bpy.types.Operator):
	"Set the image to the active node's image"
	bl_idname='zoom.get_image'
	bl_label='Get Image'

	def execute(self,context):
	nodes, links = get_nodes_links(context)
	found_any=False
	for node in nodes:
	if node.select==True and node.type=="IMAGE":
	found_any=True
	context.scene.ImageName=node.image.name
	if not found_any:
	self.report({'INFO'}, "No Image node selected!")
	return {'FINISHED'}

	class ZoomFit(bpy.types.Operator):
	'Fit the background image'
	bl_idname='zoom.fit'
	bl_label='Zoom Fit'

	@classmethod
	def poll(cls, context):
	snode = context.space_data
	return snode.tree_type == 'CompositorNodeTree'

	def execute(self,context):
	renderwidth=int(actualRes(context, "x"))
	renderheight=int(actualRes(context, "y"))
	width=100 #declare window size vars
	height=100
	sbsize=8 #Scrollbar size (set to 0 if not accounting for scrollbars), default 8
	clamp=context.scene.Clamp
	fit=context.scene.Fit
	margin=context.scene.Margin
	grace=context.scene.Grace

	if context.scene.ImageType == 'image':
	img=bpy.data.images[context.scene.ImageName]
	renderwidth=img.size[0]
	renderheight=img.size[1]
	elif context.scene.ImageType == 'custom':
	renderwidth=context.scene.ImgSizeX
	renderheight=context.scene.ImgSizeY

	for region in bpy.context.area.regions:
	if region.type=="WINDOW":
	width=region.width
	height=region.height
	for space in bpy.context.area.spaces:
	if space.type=="NODE_EDITOR":
	space.backdrop_x=sbsize*-1 #Account for scroll bars
	space.backdrop_y=sbsize

	if fit:
	ratiox=width/renderwidth
	ratioy=height/renderheight
	space.backdrop_zoom=1
	if (width-(2*sbsize)+grace<renderwidth) or (clamp==False):
	space.backdrop_zoom=ratiox-(margin/100) #Fit in window
	if height-(2sbsize)+grace<renderheightratiox:
	space.backdrop_zoom=ratioy-(margin/100)
	else:
	space.backdrop_zoom=1
	return {'FINISHED'}

	class NodeViewingPanel(bpy.types.Panel):
	bl_idname = "NODE_PT_viewing"
	bl_space_type = 'NODE_EDITOR'
	bl_region_type = 'UI'
	bl_label = "Node Wrangler"

	def draw(self, context):
	type = context.space_data.tree_type
	layout = self.layout
	scene=context.scene

	if context.space_data.tree_type=='CompositorNodeTree':
	box=layout.box()
	col=box.column(align=True)
	col.operator("zoom.fit",icon="ZOOM_SELECTED", text="Backdrop Fit")
	col.prop(scene, "Fit")
	if scene.Fit:
	col.prop(scene, "Clamp")
	col.prop(scene, "Margin")
	col.prop(scene, "Grace")
	col.separator()
	col.row().prop(scene, "ImageType", expand=True)
	if context.scene.ImageType=='custom':
	row=col.row(align=True)
	row.prop(scene, "ImgSizeX", text="X")
	row.prop(scene, "ImgSizeY", text="Y")
	elif context.scene.ImageType=='image':
	col=box.column(align=True)
	col.alignment='RIGHT'
	if scene.ImageName=="":
	col.operator("zoom.get_image",icon="EXPORT")
	else:
	dx=bpy.data.images[context.scene.ImageName].size[0]
	dy=bpy.data.images[context.scene.ImageName].size[1]
	col.prop_search(scene, "ImageName", bpy.data, 'images', text="")
	col.label(text="(Size: "+str(dx)+" x "+str(dy)+")")

	box=layout.box()
	col=box.column(align=True)
	col.operator("nodes.layout",icon="IMGDISPLAY", text="Arrange Node Tree")
	col.operator("nodes.layout_sel",icon="IMGDISPLAY", text="Arrange Selection")
	col.prop(scene, "StartAlign")
	col.prop(scene, "EndAlign")
	col.prop(scene, "DelReroutes")
	#col.prop(scene, "SelectionOnly")
	col.prop(scene, "Spacing")
	col.prop(scene, "FrameHandling")

	col=layout.column(align=True)
	#box=col.box()
	col.operator("nodes.del_unused",icon="CANCEL", text="Delete Unused Nodes")


	classes = [ZoomFit, NodeViewingPanel, ArrangeNodes, DeleteUnusedNodes, GetImage, ArrangeSelectedNodes]
	addon_keymaps = []
	def register():
	# props
	bpy.types.Scene.Clamp = bpy.props.BoolProperty(
	name="Max 1.0",
	default=True,
	description="Force the Zoom to stay below or at 1.0 even when there's space for a bigger image")
	bpy.types.Scene.Fit = bpy.props.BoolProperty(
	name="Fit",
	default=True,
	description="Scale the image to fit inside the window")
	bpy.types.Scene.Margin = bpy.props.FloatProperty(
	name="Margin Size",
	default=5.0,
	min=0.0,
	description="Create margins around image when Fitting (percentage)")
	bpy.types.Scene.Grace = bpy.props.FloatProperty(
	name="Grace",
	default=0.0,
	min=0.0,
	description="When the image is only slightly smaller than the window, set Zoom to 1.0 anyway")
	bpy.types.Scene.StartAlign = bpy.props.BoolProperty(
	name="Align Start Nodes",
	default=True,
	description="Put all nodes with no inputs on the left of the tree")
	bpy.types.Scene.EndAlign = bpy.props.BoolProperty(
	name="Align End Nodes",
	default=True,
	description="Put all nodes with no outputs on the right of the tree")
	bpy.types.Scene.Spacing = bpy.props.FloatProperty(
	name="Spacing",
	default=80.0,
	min=0.0,
	description="The horizonal space between nodes (vertical is half this)")
	bpy.types.Scene.DelReroutes = bpy.props.BoolProperty(
	name="Delete Reroutes",
	default=True,
	description="Delete all Reroute nodes to avoid unexpected layouts")
	bpy.types.Scene.ImageType = bpy.props.EnumProperty(
	name="Image Type",
	items=(("render","Render","Use the render size to calculate Zoom Fit"),("image","Image","Use a specific image's size to calculate Zoom Fit"),("custom","Custom","Use custom dimensions to calculate Zoom Fit")),
	default='render',
	description="Which dimensions to calculate Zoom Fit from")
	bpy.types.Scene.ImgSizeX = bpy.props.IntProperty(
	name="Image Size",
	default=1024,
	min=0,
	description="The number of horizontal pixels")
	bpy.types.Scene.ImgSizeY = bpy.props.IntProperty(
	name="Image Size",
	default=1024,
	min=0,
	description="The number of vertical pixels")
	bpy.types.Scene.ImageName = bpy.props.StringProperty(
	name="Image",
	default='',
	description="The image whose dimensions will be used")
	bpy.types.Scene.FrameHandling = bpy.props.EnumProperty(
	name="Frames",
	items=(("ignore","Ignore","Do nothing about Frame nodes (can be messy)"),("delete","Delete","Delete Frame nodes")),
	default='ignore',
	description="How to handle Frame nodes")
	bpy.types.Scene.SelectionOnly = bpy.props.BoolProperty(
	name="Selection Only",
	default=False,
	description="Arranges just the selected nodes, not the whole tree")

	# add operator
	for c in classes:
	bpy.utils.register_class(c)

	# add keymap entry
	km = bpy.context.window_manager.keyconfigs.addon.keymaps.new(name="Node Editor", space_type="NODE_EDITOR")
	kmi = km.keymap_items.new("zoom.fit", 'Z', 'PRESS')
	kmi = km.keymap_items.new("nodes.layout", 'Q', 'PRESS', False, False, True, False)
	kmi = km.keymap_items.new("nodes.layout_sel", 'Q', 'PRESS')
	addon_keymaps.append(km)


	def unregister():
	# props
	del bpy.types.Scene.Clamp
	del bpy.types.Scene.Fit
	del bpy.types.Scene.Margin
	del bpy.types.Scene.Grace
	del bpy.types.Scene.StartAlign
	del bpy.types.Scene.EndAlign
	del bpy.types.Scene.Spacing
	del bpy.types.Scene.DelReroutes
	del bpy.types.Scene.ImageType
	del bpy.types.Scene.ImgSizeX
	del bpy.types.Scene.ImgSizeY
	del bpy.types.Scene.ImageName
	del bpy.types.Scene.FrameHandling
	del bpy.types.Scene.SelectionOnly

	# remove operator
	for c in classes:
	bpy.utils.unregister_class(c)

	# remove keymap entry
	for km in addon_keymaps:
	bpy.context.window_manager.keyconfigs.addon.keymaps.remove(km)
	addon_keymaps.clear()

	if __name__ == "__main__":
	register()

File Metadata

Mime Type: text/x-c++
Storage Engine: local-disk
Storage Format: Raw Data
Storage Handle: 79/31/4d0002434b48609056530a4b6e03

node_wrangler.pyPublicActions

node_wrangler.pyView Options

File Metadata

Event Timeline

node_wrangler.py
Public
Actions

node_wrangler.py
View Options