From 8dbd8a8d9293556872292352a0baaff7680dba71 Mon Sep 17 00:00:00 2001 From: David Friedli Date: Fri, 24 Apr 2020 20:57:23 +0200 Subject: Node Wrangler: support emission viewer inside node groups Differential Revision: https://developer.blender.org/D6956 Reviewers: gregzaal --- node_wrangler.py | 337 +++++++++++++++++++++++++++++++++++++++++++++---------- 1 file changed, 280 insertions(+), 57 deletions(-) (limited to 'node_wrangler.py') diff --git a/node_wrangler.py b/node_wrangler.py index e26da88b..717ffd8e 100644 --- a/node_wrangler.py +++ b/node_wrangler.py @@ -596,9 +596,10 @@ draw_color_sets = { ) } +viewer_socket_name = "tmp_viewer" def is_visible_socket(socket): - return not socket.hide and socket.enabled + return not socket.hide and socket.enabled and socket.type != 'CUSTOM' def nice_hotkey_name(punc): # convert the ugly string name into the actual character @@ -1030,10 +1031,9 @@ def draw_callback_nodeoutline(self, context, mode): bgl.glDisable(bgl.GL_BLEND) bgl.glDisable(bgl.GL_LINE_SMOOTH) - -def get_nodes_links(context): +def get_active_tree(context): tree = context.space_data.node_tree - + path = [] # 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 @@ -1042,9 +1042,71 @@ def get_nodes_links(context): if tree.nodes.active: while tree.nodes.active != context.active_node: tree = tree.nodes.active.node_tree + path.append(tree) + return tree, path +def get_nodes_links(context): + tree, path = get_active_tree(context) return tree.nodes, tree.links +def is_viewer_socket(socket): + # checks if a internal socket is a valid viewer socket + return socket.name == viewer_socket_name and socket.NWViewerSocket + +def get_internal_socket(socket): + #get the internal socket from a socket inside or outside the group + node = socket.node + if node.type == 'GROUP_OUTPUT': + source_iterator = node.inputs + iterator = node.id_data.outputs + elif node.type == 'GROUP_INPUT': + source_iterator = node.outputs + iterator = node.id_data.inputs + elif hasattr(node, "node_tree"): + if socket.is_output: + source_iterator = node.outputs + iterator = node.node_tree.outputs + else: + source_iterator = node.inputs + iterator = node.node_tree.inputs + else: + return None + + for i, s in enumerate(source_iterator): + if s == socket: + break + return iterator[i] + +def is_viewer_link(link, output_node): + if "Emission Viewer" in link.to_node.name or link.to_node == output_node and link.to_socket == output_node.inputs[0]: + return True + if link.to_node.type == 'GROUP_OUTPUT': + socket = get_internal_socket(link.to_socket) + if is_viewer_socket(socket): + return True + return False + +def get_group_output_node(tree): + for node in tree.nodes: + if node.type == 'GROUP_OUTPUT' and node.is_active_output == True: + return node + +def get_output_location(tree): + # get right-most location + sorted_by_xloc = (sorted(tree.nodes, key=lambda x: x.location.x)) + max_xloc_node = sorted_by_xloc[-1] + if max_xloc_node.name == 'Emission Viewer': + max_xloc_node = sorted_by_xloc[-2] + + # get average y location + sum_yloc = 0 + for node in tree.nodes: + sum_yloc += node.location.y + + loc_x = max_xloc_node.location.x + max_xloc_node.dimensions.x + 80 + loc_y = sum_yloc / len(tree.nodes) + return loc_x, loc_y + # Principled prefs class NWPrincipledPreferences(bpy.types.PropertyGroup): base_color: StringProperty( @@ -1624,13 +1686,17 @@ class NWAddAttrNode(Operator, NWBase): nodes.active.attribute_name = self.attr_name return {'FINISHED'} - class NWEmissionViewer(Operator, NWBase): bl_idname = "node.nw_emission_viewer" bl_label = "Emission Viewer" bl_description = "Connect active node to Emission Shader for shadeless previews" bl_options = {'REGISTER', 'UNDO'} + def __init__(self): + self.shader_output_type = "" + self.shader_output_ident = "" + self.shader_viewer_ident = "" + @classmethod def poll(cls, context): if nw_check(context): @@ -1643,35 +1709,160 @@ class NWEmissionViewer(Operator, NWBase): return True return False - def invoke(self, context, event): - space = context.space_data - shader_type = space.shader_type + def ensure_viewer_socket(self, node, socket_type, connect_socket=None): + #check if a viewer output already exists in a node group otherwise create + if hasattr(node, "node_tree"): + index = None + if len(node.node_tree.outputs): + free_socket = None + for i, socket in enumerate(node.node_tree.outputs): + if is_viewer_socket(socket) and is_visible_socket(node.outputs[i]) and socket.type == socket_type: + #if viewer output is already used but leads to the same socket we can still use it + is_used = self.is_socket_used_other_mats(socket) + if is_used: + if connect_socket == None: + continue + groupout = get_group_output_node(node.node_tree) + groupout_input = groupout.inputs[i] + links = groupout_input.links + if connect_socket not in [link.from_socket for link in links]: + continue + index=i + break + if not free_socket: + free_socket = i + if not index and free_socket: + index = free_socket + + if not index: + #create viewer socket + node.node_tree.outputs.new(socket_type, viewer_socket_name) + index = len(node.node_tree.outputs) - 1 + node.node_tree.outputs[index].NWViewerSocket = True + return index + + def init_shader_variables(self, space, shader_type): if shader_type == 'OBJECT': if space.id not in [light for light in bpy.data.lights]: # cannot use bpy.data.lights directly as iterable - shader_output_type = "OUTPUT_MATERIAL" - shader_output_ident = "ShaderNodeOutputMaterial" - shader_viewer_ident = "ShaderNodeEmission" + self.shader_output_type = "OUTPUT_MATERIAL" + self.shader_output_ident = "ShaderNodeOutputMaterial" + self.shader_viewer_ident = "ShaderNodeEmission" else: - shader_output_type = "OUTPUT_LIGHT" - shader_output_ident = "ShaderNodeOutputLight" - shader_viewer_ident = "ShaderNodeEmission" + self.shader_output_type = "OUTPUT_LIGHT" + self.shader_output_ident = "ShaderNodeOutputLight" + self.shader_viewer_ident = "ShaderNodeEmission" elif shader_type == 'WORLD': - shader_output_type = "OUTPUT_WORLD" - shader_output_ident = "ShaderNodeOutputWorld" - shader_viewer_ident = "ShaderNodeBackground" + self.shader_output_type = "OUTPUT_WORLD" + self.shader_output_ident = "ShaderNodeOutputWorld" + self.shader_viewer_ident = "ShaderNodeBackground" + + def get_shader_output_node(self, tree): + for node in tree.nodes: + if node.type == self.shader_output_type and node.is_active_output == True: + return node + + @classmethod + def ensure_group_output(cls, tree): + #check if a group output node exists otherwise create + groupout = get_group_output_node(tree) + if not groupout: + groupout = tree.nodes.new('NodeGroupOutput') + loc_x, loc_y = get_output_location(tree) + groupout.location.x = loc_x + groupout.location.y = loc_y + groupout.select = False + return groupout + + @classmethod + def search_sockets(cls, node, sockets, index=None): + #recursevley scan nodes for viewer sockets and store in list + for i, input_socket in enumerate(node.inputs): + if index and i != index: + continue + if len(input_socket.links): + link = input_socket.links[0] + next_node = link.from_node + external_socket = link.from_socket + if hasattr(next_node, "node_tree"): + for socket_index, s in enumerate(next_node.outputs): + if s == external_socket: + break + socket = next_node.node_tree.outputs[socket_index] + if is_viewer_socket(socket) and socket not in sockets: + sockets.append(socket) + #continue search inside of node group but restrict socket to where we came from + groupout = get_group_output_node(next_node.node_tree) + cls.search_sockets(groupout, sockets, index=socket_index) + + @classmethod + def scan_nodes(cls, tree, selection, sockets): + # get all selcted nodes and all viewer sockets in a material tree + for node in tree.nodes: + if node.select: + selection.append(node) + node.select = False + + if hasattr(node, "node_tree"): + for socket in node.node_tree.outputs: + if is_viewer_socket(socket) and (socket not in sockets): + sockets.append(socket) + cls.scan_nodes(node.node_tree, selection, sockets) + + def link_leads_to_used_socket(self, link): + #return True if link leads to a socket that is already used in this material + socket = get_internal_socket(link.to_socket) + return (socket and self.is_socket_used_active_mat(socket)) + + def is_socket_used_active_mat(self, socket): + #ensure used sockets in active material is calculated and check given socket + if not hasattr(self, "used_viewer_sockets_active_mat"): + self.used_viewer_sockets_active_mat = [] + materialout = self.get_shader_output_node(bpy.context.space_data.node_tree) + if materialout: + emission = self.get_viewer_node(materialout) + self.search_sockets((emission if emission else materialout), self.used_viewer_sockets_active_mat) + return socket in self.used_viewer_sockets_active_mat + + def is_socket_used_other_mats(self, socket): + #ensure used sockets in other materials are calculated and check given socket + if not hasattr(self, "used_viewer_sockets_other_mats"): + self.used_viewer_sockets_other_mats = [] + for mat in bpy.data.materials: + if mat.node_tree == bpy.context.space_data.node_tree or not hasattr(mat.node_tree, "nodes"): + continue + # get viewer node + materialout = self.get_shader_output_node(mat.node_tree) + if materialout: + emission = self.get_viewer_node(materialout) + self.search_sockets((emission if emission else materialout), self.used_viewer_sockets_other_mats) + return socket in self.used_viewer_sockets_other_mats + + @staticmethod + def get_viewer_node(materialout): + input_socket = materialout.inputs[0] + if len(input_socket.links) > 0: + node = input_socket.links[0].from_node + if node.type == 'EMISSION' and node.name == "Emission Viewer": + return node + + def invoke(self, context, event): + space = context.space_data + shader_type = space.shader_type + self.init_shader_variables(space, shader_type) shader_types = [x[1] for x in shaders_shader_nodes_props] mlocx = event.mouse_region_x mlocy = event.mouse_region_y select_node = bpy.ops.node.select(mouse_x=mlocx, mouse_y=mlocy, extend=False) if 'FINISHED' in select_node: # only run if mouse click is on a node - nodes, links = get_nodes_links(context) - in_group = context.active_node != space.node_tree.nodes.active + active_tree, path_to_tree = get_active_tree(context) + nodes, links = active_tree.nodes, active_tree.links + base_node_tree = space.node_tree active = nodes.active output_types = [x[1] for x in shaders_output_nodes_props] valid = False if active: - if (active.name != "Emission Viewer") and (active.type not in output_types) and not in_group: + if (active.name != "Emission Viewer") and (active.type not in output_types): for out in active.outputs: if is_visible_socket(out): valid = True @@ -1680,30 +1871,15 @@ class NWEmissionViewer(Operator, NWBase): # get material_output node, store selection, deselect all materialout = None # placeholder node selection = [] - for node in nodes: - if node.type == shader_output_type: - materialout = node - if node.select: - selection.append(node.name) - node.select = False - if not materialout: - # get right-most location - sorted_by_xloc = (sorted(nodes, key=lambda x: x.location.x)) - max_xloc_node = sorted_by_xloc[-1] - if max_xloc_node.name == 'Emission Viewer': - max_xloc_node = sorted_by_xloc[-2] - - # get average y location - sum_yloc = 0 - for node in nodes: - sum_yloc += node.location.y + delete_sockets = [] - new_locx = max_xloc_node.location.x + max_xloc_node.dimensions.x + 80 - new_locy = sum_yloc / len(nodes) + #scan through every single node in tree including nodes inside of groups + self.scan_nodes(base_node_tree, selection, delete_sockets) - materialout = nodes.new(shader_output_ident) - materialout.location.x = new_locx - materialout.location.y = new_locy + materialout = self.get_shader_output_node(base_node_tree) + if not materialout: + materialout = base_node_tree.nodes.new(self.shader_output_ident) + materialout.location = get_output_location(base_node_tree) materialout.select = False # Analyze outputs, add "Emission Viewer" if needed, make links out_i = None @@ -1715,24 +1891,28 @@ class NWEmissionViewer(Operator, NWBase): out_i = valid_outputs[0] # Start index of node's outputs for i, valid_i in enumerate(valid_outputs): for out_link in active.outputs[valid_i].links: - if "Emission Viewer" in out_link.to_node.name or (out_link.to_node == materialout and out_link.to_socket == materialout.inputs[0]): - if i < len(valid_outputs) - 1: - out_i = valid_outputs[i + 1] - else: - out_i = valid_outputs[0] + if is_viewer_link(out_link, materialout): + if nodes == base_node_tree.nodes or self.link_leads_to_used_socket(out_link): + if i < len(valid_outputs) - 1: + out_i = valid_outputs[i + 1] + else: + out_i = valid_outputs[0] + make_links = [] # store sockets for new links + delete_nodes = [] # store unused nodes to delete in the end if active.outputs: # If output type not 'SHADER' - "Emission Viewer" needed if active.outputs[out_i].type != 'SHADER': + socket_type = 'NodeSocketColor' # get Emission Viewer node emission_exists = False - emission_placeholder = nodes[0] - for node in nodes: + emission_placeholder = base_node_tree.nodes[0] + for node in base_node_tree.nodes: if "Emission Viewer" in node.name: emission_exists = True emission_placeholder = node if not emission_exists: - emission = nodes.new(shader_viewer_ident) + emission = base_node_tree.nodes.new(self.shader_viewer_ident) emission.hide = True emission.location = [materialout.location.x, (materialout.location.y + 40)] emission.label = "Viewer" @@ -1742,7 +1922,7 @@ class NWEmissionViewer(Operator, NWBase): emission.select = False else: emission = emission_placeholder - make_links.append((active.outputs[out_i], emission.inputs[0])) + output_socket = emission.inputs[0] # If Viewer is connected to output by user, don't change those connections (patch by gandalf3) if emission.outputs[0].links.__len__() > 0: @@ -1762,17 +1942,54 @@ class NWEmissionViewer(Operator, NWBase): else: # Output type is 'SHADER', no Viewer needed. Delete Viewer if exists. - make_links.append((active.outputs[out_i], materialout.inputs[1 if active.outputs[out_i].name == "Volume" else 0])) - for node in nodes: + socket_type = 'NodeSocketShader' + materialout_index = 1 if active.outputs[out_i].name == "Volume" else 0 + make_links.append((active.outputs[out_i], materialout.inputs[materialout_index])) + output_socket = materialout.inputs[materialout_index] + for node in base_node_tree.nodes: if node.name == 'Emission Viewer': - node.select = True - bpy.ops.node.delete() + delete_nodes.append((base_node_tree, node)) for li_from, li_to in make_links: - links.new(li_from, li_to) + base_node_tree.links.new(li_from, li_to) + + # Crate links through node groups until we reach the active node + tree = base_node_tree + link_end = output_socket + while tree.nodes.active != active: + node = tree.nodes.active + index = self.ensure_viewer_socket(node, socket_type, connect_socket=active.outputs[out_i] if node.node_tree.nodes.active == active else None) + link_start = node.outputs[index] + node_socket = node.node_tree.outputs[index] + if node_socket in delete_sockets: + delete_sockets.remove(node_socket) + tree.links.new(link_start, link_end) + # Iterate + link_end = self.ensure_group_output(node.node_tree).inputs[index] + tree = tree.nodes.active.node_tree + tree.links.new(active.outputs[out_i], link_end) + + # Delete sockets + for socket in delete_sockets: + if not self.is_socket_used_other_mats(socket): + tree = socket.id_data + tree.outputs.remove(socket) + + # Delete nodes + for node in delete_nodes: + space.node_tree = node[0] + node[1].select = True + bpy.ops.node.delete() + # Restore selection + path = space.path + path.start(base_node_tree) + if len(path_to_tree): + for tree in path_to_tree: + path.append(tree) + nodes.active = active for node in nodes: - if node.name in selection: + if node in selection: node.select = True force_update(context) return {'FINISHED'} @@ -4847,6 +5064,11 @@ def register(): name="Source Socket!", default=0, description="An internal property used to store the source socket in a Lazy Connect operation") + bpy.types.NodeSocketInterface.NWViewerSocket = BoolProperty( + name="NW Socket", + default=False, + description="An internal property used to determine if a socket is generated by the addon" + ) for cls in classes: register_class(cls) @@ -4882,6 +5104,7 @@ def unregister(): del bpy.types.Scene.NWLazySource del bpy.types.Scene.NWLazyTarget del bpy.types.Scene.NWSourceSocket + del bpy.types.NodeSocketInterface.NWViewerSocket # keymaps for km, kmi in addon_keymaps: -- cgit v1.2.3