node_wrangler/__init__.py

@@ -1284,9 +1284,13 @@ class NWAddAttrNode(Operator, NWBase):
 class NWPreviewNode(Operator, NWBase):
     bl_idname = "node.nw_preview_node"
     bl_label = "Preview Node"
-    bl_description = "Connect the chosen node to the output or the Viewer Node"
+    bl_description = "Connect active node to the Node Group output or the Material Output"
     bl_options = {'REGISTER', 'UNDO'}
 
+    # If false, the operator is not executed if the current node group happens to be a geometry nodes group.
+    # This is needed because geometry nodes has its own viewer node that uses the same shortcut as in the compositor.
+    run_in_geometry_nodes: BoolProperty(default=True)
+
     def __init__(self):
         self.shader_output_type = ""
         self.shader_output_ident = ""
@@ -1295,7 +1299,7 @@ class NWPreviewNode(Operator, NWBase):
     def poll(cls, context):
         if nw_check(context):
             space = context.space_data
-            if space.tree_type == 'ShaderNodeTree':
+            if space.tree_type == 'ShaderNodeTree' or space.tree_type == 'GeometryNodeTree':
                 if context.active_node:
                     if context.active_node.type != "OUTPUT_MATERIAL" or context.active_node.type != "OUTPUT_WORLD":
                         return True
@@ -1429,6 +1433,9 @@ class NWPreviewNode(Operator, NWBase):
 
     def invoke(self, context, event):
         space = context.space_data
+        # Ignore operator when running in wrong context.
+        if self.run_in_geometry_nodes != (space.tree_type == "GeometryNodeTree"):
+            return {'PASS_THROUGH'}
 
         shader_type = space.shader_type
         self.init_shader_variables(space, shader_type)
@@ -1441,6 +1448,90 @@ class NWPreviewNode(Operator, NWBase):
             base_node_tree = space.node_tree
             active = nodes.active
 
+            # For geometry node trees we just connect to the group output
+            if space.tree_type == "GeometryNodeTree":
+                valid = False
+                if active:
+                    for out in active.outputs:
+                        if is_visible_socket(out):
+                            valid = True
+                            break
+                # Exit early
+                if not valid:
+                    return {'FINISHED'}
+
+                delete_sockets = []
+
+                # Scan through all nodes in tree including nodes inside of groups to find viewer sockets
+                self.scan_nodes(base_node_tree, delete_sockets)
+
+                # Find (or create if needed) the output of this node tree
+                geometryoutput = self.ensure_group_output(base_node_tree)
+
+                # Analyze outputs, make links
+                out_i = None
+                valid_outputs = []
+                for i, out in enumerate(active.outputs):
+                    if is_visible_socket(out) and out.type == 'GEOMETRY':
+                        valid_outputs.append(i)
+                if valid_outputs:
+                    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 is_viewer_link(out_link, geometryoutput):
+                            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
+                if active.outputs:
+                    # If there is no 'GEOMETRY' output type - We can't preview the node
+                    if out_i is None:
+                        return {'FINISHED'}
+                    socket_type = 'GEOMETRY'
+                    # Find an input socket of the output of type geometry
+                    geometryoutindex = None
+                    for i,inp in enumerate(geometryoutput.inputs):
+                        if inp.type == socket_type:
+                            geometryoutindex = i
+                            break
+                    if geometryoutindex is None:
+                        # Create geometry socket
+                        geometryoutput.inputs.new(socket_type, 'Geometry')
+                        geometryoutindex = len(geometryoutput.inputs) - 1
+
+                    make_links.append((active.outputs[out_i], geometryoutput.inputs[geometryoutindex]))
+                    output_socket = geometryoutput.inputs[geometryoutindex]
+                    for li_from, li_to in make_links:
+                        base_node_tree.links.new(li_from, li_to)
+                    tree = base_node_tree
+                    link_end = output_socket
+                    while tree.nodes.active != active:
+                        node = tree.nodes.active
+                        index = self.ensure_viewer_socket(node,'NodeSocketGeometry', 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:
+                    tree = socket.id_data
+                    tree.outputs.remove(socket)
+
+                nodes.active = active
+                active.select = True
+                force_update(context)
+                return {'FINISHED'}
+
+
             # What follows is code for the shader editor
             output_types = [x.nodetype for x in
                             get_nodes_from_category('Output', context)]
@@ -4257,7 +4348,8 @@ kmi_defs = (
     # Swap Links
     (NWSwapLinks.bl_idname, 'S', 'PRESS', False, False, True, None, "Swap Links"),
     # Preview Node
-    (NWPreviewNode.bl_idname, 'LEFTMOUSE', 'PRESS', True, True, False, None, "Preview node output"),
+    (NWPreviewNode.bl_idname, 'LEFTMOUSE', 'PRESS', True, True, False, (('run_in_geometry_nodes', False),), "Preview node output"),
+    (NWPreviewNode.bl_idname, 'LEFTMOUSE', 'PRESS', False, True, True, (('run_in_geometry_nodes', True),), "Preview node output"),
     # Reload Images
     (NWReloadImages.bl_idname, 'R', 'PRESS', False, False, True, None, "Reload images"),
     # Lazy Mix