Node Wrangler: support emission viewer inside node groups

Differential Revision: https://developer.blender.org/D6956

Reviewers: gregzaal

1 files changed, 280 insertions, 57 deletions

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: