diff --git a/node_wrangler/__init__.py b/node_wrangler/__init__.py index e6e6fef6c..a127ba8c0 100644 --- a/node_wrangler/__init__.py +++ b/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