Fix T92871: Node Wrangler: connect to output shortcut in Geo Node

This implements the use of shortcut `O` to connect the first geometry
output of the active node, to the group output.

Right now, only geometry is concerned, and the operator is cancelled
if another output type is found, such as an `int`.

1 files changed, 14 insertions, 6 deletions

diff --git a/node_wrangler.py b/node_wrangler.py
index 7e1e932e..30331ce2 100644
--- a/node_wrangler.py
+++ b/node_wrangler.py
@@ -3748,7 +3748,7 @@ class NWLinkToOutputNode(Operator):
     @classmethod
     def poll(cls, context):
         valid = False
-        if nw_check(context) and context.space_data.tree_type != 'GeometryNodeTree':
+        if nw_check(context):
             if context.active_node is not None:
                 for out in context.active_node.outputs:
                     if is_visible_socket(out):
@@ -3762,11 +3762,14 @@ class NWLinkToOutputNode(Operator):
         output_node = None
         output_index = None
         tree_type = context.space_data.tree_type
-        output_types_shaders = [x[1] for x in shaders_output_nodes_props]
-        output_types_compo = ['COMPOSITE']
-        output_types_blender_mat = ['OUTPUT']
-        output_types_textures = ['OUTPUT']
-        output_types = output_types_shaders + output_types_compo + output_types_blender_mat
+        if tree_type == 'ShaderNodeTree':
+            output_types = [x[1] for x in shaders_output_nodes_props] + ['OUTPUT']
+        elif tree_type == 'CompositorNodeTree':
+            output_types = ['COMPOSITE']
+        elif tree_type == 'TextureNodeTree':
+            output_types = ['OUTPUT']
+        elif tree_type == 'GeometryNodeTree':
+            output_types = ['GROUP_OUTPUT']
         for node in nodes:
             if node.type in output_types:
                 output_node = node
@@ -3779,6 +3782,8 @@ class NWLinkToOutputNode(Operator):
                 output_node = nodes.new('CompositorNodeComposite')
             elif tree_type == 'TextureNodeTree':
                 output_node = nodes.new('TextureNodeOutput')
+            elif tree_type == 'GeometryNodeTree':
+                output_node = nodes.new('NodeGroupOutput')
             output_node.location.x = active.location.x + active.dimensions.x + 80
             output_node.location.y = active.location.y
         if (output_node and active.outputs):
@@ -3797,6 +3802,9 @@ class NWLinkToOutputNode(Operator):
                     out_input_index = 1
                 elif active.outputs[output_index].type != 'SHADER':  # connect to displacement if not a shader
                     out_input_index = 2
+            elif tree_type == 'GeometryNodeTree':
+                if active.outputs[output_index].type != 'GEOMETRY':
+                    return {'CANCELLED'}
             links.new(active.outputs[output_index], output_node.inputs[out_input_index])
 
         force_update(context)  # viewport render does not update