Fix T96194: Node Wrangler: use enabled sockets for lazy mix

Add utility function to check for the first output that is actually enabled
rather than always using the first output of a node.

The previous behavior would lead to confusing situations, where a mix/math
node was added but didn't seem to be connected, when the first output was
currently disabled.

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

1 files changed, 21 insertions, 12 deletions

diff --git a/node_wrangler.py b/node_wrangler.py
index 02aa62c4..eed4c2fb 100644
--- a/node_wrangler.py
+++ b/node_wrangler.py
@@ -596,6 +596,13 @@ def get_nodes_from_category(category_name, context):
         if category.name == category_name:
             return sorted(category.items(context), key=lambda node: node.label)
 
+def get_first_enabled_output(node):
+    for output in node.outputs:
+        if output.enabled:
+            return output
+    else:
+        return node.outputs[0]
+
 def is_visible_socket(socket):
     return not socket.hide and socket.enabled and socket.type != 'CUSTOM'
 
@@ -2452,8 +2459,8 @@ class NWMergeNodes(Operator, NWBase):
             mode = 'MIX'
         if (merge_type != 'MATH' and merge_type != 'GEOMETRY') and tree_type == 'GEOMETRY':
             merge_type = 'AUTO'
-        # The math nodes used for geometry nodes are of type 'ShaderNode'
-        if merge_type == 'MATH' and tree_type == 'GEOMETRY':
+        # The MixRGB node and math nodes used for geometry nodes are of type 'ShaderNode'
+        if (merge_type == 'MATH' or merge_type == 'MIX') and tree_type == 'GEOMETRY':
             node_type = 'ShaderNode'
         selected_mix = []  # entry = [index, loc]
         selected_shader = []  # entry = [index, loc]
@@ -2473,7 +2480,8 @@ class NWMergeNodes(Operator, NWBase):
                             ('VALUE', [t[0] for t in operations], selected_math),
                             ('VECTOR', [], selected_vector),
                     ):
-                        output_type = node.outputs[0].type
+                        output = get_first_enabled_output(node)
+                        output_type = output.type
                         valid_mode = mode in types_list
                         # When mode is 'MIX' we have to cheat since the mix node is not used in
                         # geometry nodes.
@@ -2524,7 +2532,7 @@ class NWMergeNodes(Operator, NWBase):
 
             # Change the node type for math nodes in a geometry node tree.
             if tree_type == 'GEOMETRY':
-                if nodes_list is selected_math or nodes_list is selected_vector:
+                if nodes_list is selected_math or nodes_list is selected_vector or nodes_list is selected_mix:
                     node_type = 'ShaderNode'
                     if mode == 'MIX':
                         mode = 'ADD'
@@ -2649,8 +2657,9 @@ class NWMergeNodes(Operator, NWBase):
             # Special case:
             # Two nodes were selected and first selected has no output links, second selected has output links.
             # Then add links from last add to all links 'to_socket' of out links of second selected.
+            first_selected_output = get_first_enabled_output(first_selected)
             if len(nodes_list) == 2:
-                if not first_selected.outputs[0].links:
+                if not first_selected_output.links:
                     second_selected = nodes[nodes_list[1][0]]
                     for ss_link in second_selected.outputs[0].links:
                         # Prevent cyclic dependencies when nodes to be merged are linked to one another.
@@ -2658,16 +2667,16 @@ class NWMergeNodes(Operator, NWBase):
                         if not self.link_creates_cycle(ss_link, invalid_nodes):
                             links.new(last_add.outputs[0], ss_link.to_socket)
             # add links from last_add to all links 'to_socket' of out links of first selected.
-            for fs_link in first_selected.outputs[0].links:
+            for fs_link in first_selected_output.links:
                 # Link only if "to_node" index not in invalid indexes list.
                 if not self.link_creates_cycle(fs_link, invalid_nodes):
                     links.new(last_add.outputs[0], fs_link.to_socket)
             # add link from "first" selected and "first" add node
             node_to = nodes[count_after - 1]
-            links.new(first_selected.outputs[0], node_to.inputs[first])
+            links.new(first_selected_output, node_to.inputs[first])
             if node_to.type == 'ZCOMBINE':
                 for fs_out in first_selected.outputs:
-                    if fs_out != first_selected.outputs[0] and fs_out.name in ('Z', 'Depth'):
+                    if fs_out != first_selected_output and fs_out.name in ('Z', 'Depth'):
                         links.new(fs_out, node_to.inputs[1])
                         break
             # add links between added ADD nodes and between selected and ADD nodes
@@ -2677,20 +2686,20 @@ class NWMergeNodes(Operator, NWBase):
                     node_to = nodes[index - 1]
                     node_to_input_i = first
                     node_to_z_i = 1  # if z combine - link z to first z input
-                    links.new(node_from.outputs[0], node_to.inputs[node_to_input_i])
+                    links.new(get_first_enabled_output(node_from), node_to.inputs[node_to_input_i])
                     if node_to.type == 'ZCOMBINE':
                         for from_out in node_from.outputs:
-                            if from_out != node_from.outputs[0] and from_out.name in ('Z', 'Depth'):
+                            if from_out != get_first_enabled_output(node_from) and from_out.name in ('Z', 'Depth'):
                                 links.new(from_out, node_to.inputs[node_to_z_i])
                 if len(nodes_list) > 1:
                     node_from = nodes[nodes_list[i + 1][0]]
                     node_to = nodes[index]
                     node_to_input_i = second
                     node_to_z_i = 3  # if z combine - link z to second z input
-                    links.new(node_from.outputs[0], node_to.inputs[node_to_input_i])
+                    links.new(get_first_enabled_output(node_from), node_to.inputs[node_to_input_i])
                     if node_to.type == 'ZCOMBINE':
                         for from_out in node_from.outputs:
-                            if from_out != node_from.outputs[0] and from_out.name in ('Z', 'Depth'):
+                            if from_out != get_first_enabled_output(node_from) and from_out.name in ('Z', 'Depth'):
                                 links.new(from_out, node_to.inputs[node_to_z_i])
                 index -= 1
             # set "last" of added nodes as active