operators.py

# SPDX-License-Identifier: GPL-2.0-or-later

import bpy

from bpy.types import Operator
from bpy.props import (
    FloatProperty,
    EnumProperty,
    BoolProperty,
    IntProperty,
    StringProperty,
    FloatVectorProperty,
    CollectionProperty,
)
from bpy_extras.io_utils import ImportHelper, ExportHelper
from bpy_extras.node_utils import connect_sockets
from mathutils import Vector
from os import path
from glob import glob
from copy import copy
from itertools import chain

from .interface import NWConnectionListInputs, NWConnectionListOutputs

from .utils.constants import blend_types, geo_combine_operations, operations, navs, get_nodes_from_category, rl_outputs
from .utils.draw import draw_callback_nodeoutline
from .utils.paths import match_files_to_socket_names, split_into_components
from .utils.nodes import (node_mid_pt, autolink, node_at_pos, get_active_tree, get_nodes_links, is_viewer_socket,
                          is_viewer_link, get_group_output_node, get_output_location, force_update, get_internal_socket,
                          nw_check, NWBase, get_first_enabled_output, is_visible_socket, viewer_socket_name)


class NWLazyMix(Operator, NWBase):
    """Add a Mix RGB/Shader node by interactively drawing lines between nodes"""
    bl_idname = "node.nw_lazy_mix"
    bl_label = "Mix Nodes"
    bl_options = {'REGISTER', 'UNDO'}

    def modal(self, context, event):
        context.area.tag_redraw()
        nodes, links = get_nodes_links(context)
        cont = True

        start_pos = [event.mouse_region_x, event.mouse_region_y]

        node1 = None
        if not context.scene.NWBusyDrawing:
            node1 = node_at_pos(nodes, context, event)
            if node1:
                context.scene.NWBusyDrawing = node1.name
        else:
            if context.scene.NWBusyDrawing != 'STOP':
                node1 = nodes[context.scene.NWBusyDrawing]

        context.scene.NWLazySource = node1.name
        context.scene.NWLazyTarget = node_at_pos(nodes, context, event).name

        if event.type == 'MOUSEMOVE':
            self.mouse_path.append((event.mouse_region_x, event.mouse_region_y))

        elif event.type == 'RIGHTMOUSE' and event.value == 'RELEASE':
            end_pos = [event.mouse_region_x, event.mouse_region_y]
            bpy.types.SpaceNodeEditor.draw_handler_remove(self._handle, 'WINDOW')

            node2 = None
            node2 = node_at_pos(nodes, context, event)
            if node2:
                context.scene.NWBusyDrawing = node2.name

            if node1 == node2:
                cont = False

            if cont:
                if node1 and node2:
                    for node in nodes:
                        node.select = False
                    node1.select = True
                    node2.select = True

                    bpy.ops.node.nw_merge_nodes(mode="MIX", merge_type="AUTO")

            context.scene.NWBusyDrawing = ""
            return {'FINISHED'}

        elif event.type == 'ESC':
            print('cancelled')
            bpy.types.SpaceNodeEditor.draw_handler_remove(self._handle, 'WINDOW')
            return {'CANCELLED'}

        return {'RUNNING_MODAL'}

    def invoke(self, context, event):
        if context.area.type == 'NODE_EDITOR':
            # the arguments we pass the the callback
            args = (self, context, 'MIX')
            # Add the region OpenGL drawing callback
            # draw in view space with 'POST_VIEW' and 'PRE_VIEW'
            self._handle = bpy.types.SpaceNodeEditor.draw_handler_add(
                draw_callback_nodeoutline, args, 'WINDOW', 'POST_PIXEL')

            self.mouse_path = []

            context.window_manager.modal_handler_add(self)
            return {'RUNNING_MODAL'}
        else:
            self.report({'WARNING'}, "View3D not found, cannot run operator")
            return {'CANCELLED'}


class NWLazyConnect(Operator, NWBase):
    """Connect two nodes without clicking a specific socket (automatically determined"""
    bl_idname = "node.nw_lazy_connect"
    bl_label = "Lazy Connect"
    bl_options = {'REGISTER', 'UNDO'}
    with_menu: BoolProperty()

    def modal(self, context, event):
        context.area.tag_redraw()
        nodes, links = get_nodes_links(context)
        cont = True

        start_pos = [event.mouse_region_x, event.mouse_region_y]

        node1 = None
        if not context.scene.NWBusyDrawing:
            node1 = node_at_pos(nodes, context, event)
            if node1:
                context.scene.NWBusyDrawing = node1.name
        else:
            if context.scene.NWBusyDrawing != 'STOP':
                node1 = nodes[context.scene.NWBusyDrawing]

        context.scene.NWLazySource = node1.name
        context.scene.NWLazyTarget = node_at_pos(nodes, context, event).name

        if event.type == 'MOUSEMOVE':
            self.mouse_path.append((event.mouse_region_x, event.mouse_region_y))

        elif event.type == 'RIGHTMOUSE' and event.value == 'RELEASE':
            end_pos = [event.mouse_region_x, event.mouse_region_y]
            bpy.types.SpaceNodeEditor.draw_handler_remove(self._handle, 'WINDOW')

            node2 = None
            node2 = node_at_pos(nodes, context, event)
            if node2:
                context.scene.NWBusyDrawing = node2.name

            if node1 == node2:
                cont = False

            link_success = False
            if cont:
                if node1 and node2:
                    original_sel = []
                    original_unsel = []
                    for node in nodes:
                        if node.select:
                            node.select = False
                            original_sel.append(node)
                        else:
                            original_unsel.append(node)
                    node1.select = True
                    node2.select = True

                    # link_success = autolink(node1, node2, links)
                    if self.with_menu:
                        if len(node1.outputs) > 1 and node2.inputs:
                            bpy.ops.wm.call_menu("INVOKE_DEFAULT", name=NWConnectionListOutputs.bl_idname)
                        elif len(node1.outputs) == 1:
                            bpy.ops.node.nw_call_inputs_menu(from_socket=0)
                    else:
                        link_success = autolink(node1, node2, links)

                    for node in original_sel:
                        node.select = True
                    for node in original_unsel:
                        node.select = False

            if link_success:
                force_update(context)
            context.scene.NWBusyDrawing = ""
            return {'FINISHED'}

        elif event.type == 'ESC':
            bpy.types.SpaceNodeEditor.draw_handler_remove(self._handle, 'WINDOW')
            return {'CANCELLED'}

        return {'RUNNING_MODAL'}

    def invoke(self, context, event):
        if context.area.type == 'NODE_EDITOR':
            nodes, links = get_nodes_links(context)
            node = node_at_pos(nodes, context, event)
            if node:
                context.scene.NWBusyDrawing = node.name

            # the arguments we pass the the callback
            mode = "LINK"
            if self.with_menu:
                mode = "LINKMENU"
            args = (self, context, mode)
            # Add the region OpenGL drawing callback
            # draw in view space with 'POST_VIEW' and 'PRE_VIEW'
            self._handle = bpy.types.SpaceNodeEditor.draw_handler_add(
                draw_callback_nodeoutline, args, 'WINDOW', 'POST_PIXEL')

            self.mouse_path = []

            context.window_manager.modal_handler_add(self)
            return {'RUNNING_MODAL'}
        else:
            self.report({'WARNING'}, "View3D not found, cannot run operator")
            return {'CANCELLED'}


class NWDeleteUnused(Operator, NWBase):
    """Delete all nodes whose output is not used"""
    bl_idname = 'node.nw_del_unused'
    bl_label = 'Delete Unused Nodes'
    bl_options = {'REGISTER', 'UNDO'}

    delete_muted: BoolProperty(
        name="Delete Muted",
        description="Delete (but reconnect, like Ctrl-X) all muted nodes",
        default=True)
    delete_frames: BoolProperty(
        name="Delete Empty Frames",
        description="Delete all frames that have no nodes inside them",
        default=True)

    def is_unused_node(self, node):
        end_types = ['OUTPUT_MATERIAL', 'OUTPUT', 'VIEWER', 'COMPOSITE',
                     'SPLITVIEWER', 'OUTPUT_FILE', 'LEVELS', 'OUTPUT_LIGHT',
                     'OUTPUT_WORLD', 'GROUP_INPUT', 'GROUP_OUTPUT', 'FRAME']
        if node.type in end_types:
            return False

        for output in node.outputs:
            if output.links:
                return False
        return True

    @classmethod
    def poll(cls, context):
        valid = False
        if nw_check(context):
            if context.space_data.node_tree.nodes:
                valid = True
        return valid

    def execute(self, context):
        nodes, links = get_nodes_links(context)

        # Store selection
        selection = []
        for node in nodes:
            if node.select:
                selection.append(node.name)

        for node in nodes:
            node.select = False

        deleted_nodes = []
        temp_deleted_nodes = []
        del_unused_iterations = len(nodes)
        for it in range(0, del_unused_iterations):
            temp_deleted_nodes = list(deleted_nodes)  # keep record of last iteration
            for node in nodes:
                if self.is_unused_node(node):
                    node.select = True
                    deleted_nodes.append(node.name)
                    bpy.ops.node.delete()

            if temp_deleted_nodes == deleted_nodes:  # stop iterations when there are no more nodes to be deleted
                break

        if self.delete_frames:
            repeat = True
            while repeat:
                frames_in_use = []
                frames = []
                repeat = False
                for node in nodes:
                    if node.parent:
                        frames_in_use.append(node.parent)
                for node in nodes:
                    if node.type == 'FRAME' and node not in frames_in_use:
                        frames.append(node)
                        if node.parent:
                            repeat = True  # repeat for nested frames
                for node in frames:
                    if node not in frames_in_use:
                        node.select = True
                        deleted_nodes.append(node.name)
                bpy.ops.node.delete()

        if self.delete_muted:
            for node in nodes:
                if node.mute:
                    node.select = True
                    deleted_nodes.append(node.name)
            bpy.ops.node.delete_reconnect()

        # get unique list of deleted nodes (iterations would count the same node more than once)
        deleted_nodes = list(set(deleted_nodes))
        for n in deleted_nodes:
            self.report({'INFO'}, "Node " + n + " deleted")
        num_deleted = len(deleted_nodes)
        n = ' node'
        if num_deleted > 1:
            n += 's'
        if num_deleted:
            self.report({'INFO'}, "Deleted " + str(num_deleted) + n)
        else:
            self.report({'INFO'}, "Nothing deleted")

        # Restore selection
        nodes, links = get_nodes_links(context)
        for node in nodes:
            if node.name in selection:
                node.select = True
        return {'FINISHED'}

    def invoke(self, context, event):
        return context.window_manager.invoke_confirm(self, event)


class NWSwapLinks(Operator, NWBase):
    """Swap the output connections of the two selected nodes, or two similar inputs of a single node"""
    bl_idname = 'node.nw_swap_links'
    bl_label = 'Swap Links'
    bl_options = {'REGISTER', 'UNDO'}

    @classmethod
    def poll(cls, context):
        valid = False
        if nw_check(context):
            if context.selected_nodes:
                valid = len(context.selected_nodes) <= 2
        return valid

    def execute(self, context):
        nodes, links = get_nodes_links(context)
        selected_nodes = context.selected_nodes
        n1 = selected_nodes[0]

        # Swap outputs
        if len(selected_nodes) == 2:
            n2 = selected_nodes[1]
            if n1.outputs and n2.outputs:
                n1_outputs = []
                n2_outputs = []

                out_index = 0
                for output in n1.outputs:
                    if output.links:
                        for link in output.links:
                            n1_outputs.append([out_index, link.to_socket])
                            links.remove(link)
                    out_index += 1

                out_index = 0
                for output in n2.outputs:
                    if output.links:
                        for link in output.links:
                            n2_outputs.append([out_index, link.to_socket])
                            links.remove(link)
                    out_index += 1

                for connection in n1_outputs:
                    try:
                        connect_sockets(n2.outputs[connection[0]], connection[1])
                    except:
                        self.report({'WARNING'},
                                    "Some connections have been lost due to differing numbers of output sockets")
                for connection in n2_outputs:
                    try:
                        connect_sockets(n1.outputs[connection[0]], connection[1])
                    except:
                        self.report({'WARNING'},
                                    "Some connections have been lost due to differing numbers of output sockets")
            else:
                if n1.outputs or n2.outputs:
                    self.report({'WARNING'}, "One of the nodes has no outputs!")
                else:
                    self.report({'WARNING'}, "Neither of the nodes have outputs!")

        # Swap Inputs
        elif len(selected_nodes) == 1:
            if n1.inputs and n1.inputs[0].is_multi_input:
                self.report({'WARNING'}, "Can't swap inputs of a multi input socket!")
                return {'FINISHED'}
            if n1.inputs:
                types = []
                i = 0
                for i1 in n1.inputs:
                    if i1.is_linked and not i1.is_multi_input:
                        similar_types = 0
                        for i2 in n1.inputs:
                            if i1.type == i2.type and i2.is_linked and not i2.is_multi_input:
                                similar_types += 1
                        types.append([i1, similar_types, i])
                    i += 1
                types.sort(key=lambda k: k[1], reverse=True)

                if types:
                    t = types[0]
                    if t[1] == 2:
                        for i2 in n1.inputs:
                            if t[0].type == i2.type == t[0].type and t[0] != i2 and i2.is_linked:
                                pair = [t[0], i2]
                        i1f = pair[0].links[0].from_socket
                        i1t = pair[0].links[0].to_socket
                        i2f = pair[1].links[0].from_socket
                        i2t = pair[1].links[0].to_socket
                        connect_sockets(i1f, i2t)
                        connect_sockets(i2f, i1t)
                    if t[1] == 1:
                        if len(types) == 1:
                            fs = t[0].links[0].from_socket
                            i = t[2]
                            links.remove(t[0].links[0])
                            if i + 1 == len(n1.inputs):
                                i = -1
                            i += 1
                            while n1.inputs[i].is_linked:
                                i += 1
                            connect_sockets(fs, n1.inputs[i])
                        elif len(types) == 2:
                            i1f = types[0][0].links[0].from_socket
                            i1t = types[0][0].links[0].to_socket
                            i2f = types[1][0].links[0].from_socket
                            i2t = types[1][0].links[0].to_socket
                            connect_sockets(i1f, i2t)
                            connect_sockets(i2f, i1t)

                else:
                    self.report({'WARNING'}, "This node has no input connections to swap!")
            else:
                self.report({'WARNING'}, "This node has no inputs to swap!")

        force_update(context)
        return {'FINISHED'}


class NWResetBG(Operator, NWBase):
    """Reset the zoom and position of the background image"""
    bl_idname = 'node.nw_bg_reset'
    bl_label = 'Reset Backdrop'
    bl_options = {'REGISTER', 'UNDO'}

    @classmethod
    def poll(cls, context):
        valid = False
        if nw_check(context):
            snode = context.space_data
            valid = snode.tree_type == 'CompositorNodeTree'
        return valid

    def execute(self, context):
        context.space_data.backdrop_zoom = 1
        context.space_data.backdrop_offset[0] = 0
        context.space_data.backdrop_offset[1] = 0
        return {'FINISHED'}


class NWAddAttrNode(Operator, NWBase):
    """Add an Attribute node with this name"""
    bl_idname = 'node.nw_add_attr_node'
    bl_label = 'Add UV map'
    bl_options = {'REGISTER', 'UNDO'}

    attr_name: StringProperty()

    def execute(self, context):
        bpy.ops.node.add_node('INVOKE_DEFAULT', use_transform=True, type="ShaderNodeAttribute")
        nodes, links = get_nodes_links(context)
        nodes.active.attribute_name = self.attr_name
        return {'FINISHED'}


class NWPreviewNode(Operator, NWBase):
    bl_idname = "node.nw_preview_node"
    bl_label = "Preview 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 = ""

    @classmethod
    def poll(cls, context):
        if nw_check(context):
            space = context.space_data
            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
                else:
                    return True
        return False

    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 is 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
                self.shader_output_type = "OUTPUT_MATERIAL"
                self.shader_output_ident = "ShaderNodeOutputMaterial"
            else:
                self.shader_output_type = "OUTPUT_LIGHT"
                self.shader_output_ident = "ShaderNodeOutputLight"

        elif shader_type == 'WORLD':
            self.shader_output_type = "OUTPUT_WORLD"
            self.shader_output_ident = "ShaderNodeOutputWorld"

    def get_shader_output_node(self, tree):
        for node in tree.nodes:
            if node.type == self.shader_output_type and node.is_active_output:
                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
            # So that we don't keep on adding new group outputs
            groupout.is_active_output = True
        return groupout

    @classmethod
    def search_sockets(cls, node, sockets, index=None):
        # recursively 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, sockets):
        # get all viewer sockets in a material tree
        for node in tree.nodes:
            if hasattr(node, "node_tree"):
                if node.node_tree is None:
                    continue
                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, 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:
                self.search_sockets(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:
                    self.search_sockets(materialout, self.used_viewer_sockets_other_mats)
        return socket in self.used_viewer_sockets_other_mats

    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)
        mlocx = event.mouse_region_x
        mlocy = event.mouse_region_y
        select_node = bpy.ops.node.select(location=(mlocx, mlocy), extend=False)
        if 'FINISHED' in select_node:  # only run if mouse click is on a node
            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

            # 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:
                        connect_sockets(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)
                        connect_sockets(link_start, link_end)
                        # Iterate
                        link_end = self.ensure_group_output(node.node_tree).inputs[index]
                        tree = tree.nodes.active.node_tree
                    connect_sockets(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)]
            valid = False
            if active:
                if active.rna_type.identifier not in output_types:
                    for out in active.outputs:
                        if is_visible_socket(out):
                            valid = True
                            break
            if valid:
                # get material_output node
                materialout = None  # placeholder node
                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)

                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
                out_i = None
                valid_outputs = []
                for i, out in enumerate(active.outputs):
                    if is_visible_socket(out):
                        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, 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
                if active.outputs:
                    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 li_from, li_to in make_links:
                        connect_sockets(li_from, li_to)

                    # Create 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)
                        connect_sockets(link_start, link_end)
                        # Iterate
                        link_end = self.ensure_group_output(node.node_tree).inputs[index]
                        tree = tree.nodes.active.node_tree
                    connect_sockets(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)

                nodes.active = active
                active.select = True

                force_update(context)

            return {'FINISHED'}
        else:
            return {'CANCELLED'}


class NWFrameSelected(Operator, NWBase):
    bl_idname = "node.nw_frame_selected"
    bl_label = "Frame Selected"
    bl_description = "Add a frame node and parent the selected nodes to it"
    bl_options = {'REGISTER', 'UNDO'}

    label_prop: StringProperty(
        name='Label',
        description='The visual name of the frame node',
        default=' '
    )
    use_custom_color_prop: BoolProperty(
        name="Custom Color",
        description="Use custom color for the frame node",
        default=False
    )
    color_prop: FloatVectorProperty(
        name="Color",
        description="The color of the frame node",
        default=(0.604, 0.604, 0.604),
        min=0, max=1, step=1, precision=3,
        subtype='COLOR_GAMMA', size=3
    )

    def draw(self, context):
        layout = self.layout
        layout.prop(self, 'label_prop')
        layout.prop(self, 'use_custom_color_prop')
        col = layout.column()
        col.active = self.use_custom_color_prop
        col.prop(self, 'color_prop', text="")

    def execute(self, context):
        nodes, links = get_nodes_links(context)
        selected = []
        for node in nodes:
            if node.select:
                selected.append(node)

        bpy.ops.node.add_node(type='NodeFrame')
        frm = nodes.active
        frm.label = self.label_prop
        frm.use_custom_color = self.use_custom_color_prop
        frm.color = self.color_prop

        for node in selected:
            node.parent = frm

        return {'FINISHED'}


class NWReloadImages(Operator):
    bl_idname = "node.nw_reload_images"
    bl_label = "Reload Images"
    bl_description = "Update all the image nodes to match their files on disk"

    @classmethod
    def poll(cls, context):
        valid = False
        if nw_check(context) and context.space_data.tree_type != 'GeometryNodeTree':
            if context.active_node is not None:
                for out in context.active_node.outputs:
                    if is_visible_socket(out):
                        valid = True
                        break
        return valid

    def execute(self, context):
        nodes, links = get_nodes_links(context)
        image_types = ["IMAGE", "TEX_IMAGE", "TEX_ENVIRONMENT", "TEXTURE"]
        num_reloaded = 0
        for node in nodes:
            if node.type in image_types:
                if node.type == "TEXTURE":
                    if node.texture:  # node has texture assigned
                        if node.texture.type in ['IMAGE', 'ENVIRONMENT_MAP']:
                            if node.texture.image:  # texture has image assigned
                                node.texture.image.reload()
                                num_reloaded += 1
                else:
                    if node.image:
                        node.image.reload()
                        num_reloaded += 1

        if num_reloaded:
            self.report({'INFO'}, "Reloaded images")
            print("Reloaded " + str(num_reloaded) + " images")
            force_update(context)
            return {'FINISHED'}
        else:
            self.report({'WARNING'}, "No images found to reload in this node tree")
            return {'CANCELLED'}


class NWSwitchNodeType(Operator, NWBase):
    """Switch type of selected nodes """
    bl_idname = "node.nw_swtch_node_type"
    bl_label = "Switch Node Type"
    bl_options = {'REGISTER', 'UNDO'}

    to_type: StringProperty(
        name="Switch to type",
        default='',
    )

    def execute(self, context):
        to_type = self.to_type
        if len(to_type) == 0:
            return {'CANCELLED'}

        nodes, links = get_nodes_links(context)
        # Those types of nodes will not swap.
        src_excludes = ('NodeFrame')
        # Those attributes of nodes will be copied if possible
        attrs_to_pass = ('color', 'hide', 'label', 'mute', 'parent',
                         'show_options', 'show_preview', 'show_texture',
                         'use_alpha', 'use_clamp', 'use_custom_color', 'location'
                         )
        selected = [n for n in nodes if n.select]
        reselect = []
        for node in [n for n in selected if
                     n.rna_type.identifier not in src_excludes and
                     n.rna_type.identifier != to_type]:
            new_node = nodes.new(to_type)
            for attr in attrs_to_pass:
                if hasattr(node, attr) and hasattr(new_node, attr):
                    setattr(new_node, attr, getattr(node, attr))
            # set image datablock of dst to image of src
            if hasattr(node, 'image') and hasattr(new_node, 'image'):
                if node.image:
                    new_node.image = node.image
            # Special cases
            if new_node.type == 'SWITCH':
                new_node.hide = True
            # Dictionaries: src_sockets and dst_sockets:
            # 'INPUTS': input sockets ordered by type (entry 'MAIN' main type of inputs).
            # 'OUTPUTS': output sockets ordered by type (entry 'MAIN' main type of outputs).
            # in 'INPUTS' and 'OUTPUTS':
            # 'SHADER', 'RGBA', 'VECTOR', 'VALUE' - sockets of those types.
            # socket entry:
            # (index_in_type, socket_index, socket_name, socket_default_value, socket_links)
            src_sockets = {
                'INPUTS': {'SHADER': [], 'RGBA': [], 'VECTOR': [], 'VALUE': [], 'MAIN': None},
                'OUTPUTS': {'SHADER': [], 'RGBA': [], 'VECTOR': [], 'VALUE': [], 'MAIN': None},
            }
            dst_sockets = {
                'INPUTS': {'SHADER': [], 'RGBA': [], 'VECTOR': [], 'VALUE': [], 'MAIN': None},
                'OUTPUTS': {'SHADER': [], 'RGBA': [], 'VECTOR': [], 'VALUE': [], 'MAIN': None},
            }
            types_order_one = 'SHADER', 'RGBA', 'VECTOR', 'VALUE'
            types_order_two = 'SHADER', 'VECTOR', 'RGBA', 'VALUE'
            # check src node to set src_sockets values and dst node to set dst_sockets dict values
            for sockets, nd in ((src_sockets, node), (dst_sockets, new_node)):
                # Check node's inputs and outputs and fill proper entries in "sockets" dict
                for in_out, in_out_name in ((nd.inputs, 'INPUTS'), (nd.outputs, 'OUTPUTS')):
                    # enumerate in inputs, then in outputs
                    # find name, default value and links of socket
                    for i, socket in enumerate(in_out):
                        the_name = socket.name
                        dval = None
                        # Not every socket, especially in outputs has "default_value"
                        if hasattr(socket, 'default_value'):
                            dval = socket.default_value
                        socket_links = []
                        for lnk in socket.links:
                            socket_links.append(lnk)
                        # check type of socket to fill proper keys.
                        for the_type in types_order_one:
                            if socket.type == the_type:
                                # create values for sockets['INPUTS'][the_type] and sockets['OUTPUTS'][the_type]
                                # entry structure: (index_in_type, socket_index, socket_name,
                                # socket_default_value, socket_links)
                                sockets[in_out_name][the_type].append(
                                    (len(sockets[in_out_name][the_type]), i, the_name, dval, socket_links))
                    # Check which of the types in inputs/outputs is considered to be "main".
                    # Set values of sockets['INPUTS']['MAIN'] and sockets['OUTPUTS']['MAIN']
                    for type_check in types_order_one:
                        if sockets[in_out_name][type_check]:
                            sockets[in_out_name]['MAIN'] = type_check
                            break

            matches = {
                'INPUTS': {'SHADER': [], 'RGBA': [], 'VECTOR': [], 'VALUE_NAME': [], 'VALUE': [], 'MAIN': []},
                'OUTPUTS': {'SHADER': [], 'RGBA': [], 'VECTOR': [], 'VALUE_NAME': [], 'VALUE': [], 'MAIN': []},
            }