Skip to content
Snippets Groups Projects
node_wrangler.py 226 KiB
Newer Older
  • Learn to ignore specific revisions
  •                         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
    
                    delete_nodes = [] # store unused nodes to delete in the end
    
                    if active.outputs:
                        # If output type not 'SHADER' - "Emission Viewer" needed
                        if active.outputs[out_i].type != 'SHADER':
    
                            socket_type = 'NodeSocketColor'
    
                            # get Emission Viewer node
                            emission_exists = False
    
                            emission_placeholder = base_node_tree.nodes[0]
                            for node in base_node_tree.nodes:
    
                                if "Emission Viewer" in node.name:
                                    emission_exists = True
                                    emission_placeholder = node
                            if not emission_exists:
    
                                emission = base_node_tree.nodes.new(self.shader_viewer_ident)
    
                                emission.hide = True
                                emission.location = [materialout.location.x, (materialout.location.y + 40)]
                                emission.label = "Viewer"
                                emission.name = "Emission Viewer"
                                emission.use_custom_color = True
                                emission.color = (0.6, 0.5, 0.4)
                                emission.select = False
                            else:
                                emission = emission_placeholder
    
                            output_socket = emission.inputs[0]
    
    
                            # If Viewer is connected to output by user, don't change those connections (patch by gandalf3)
                            if emission.outputs[0].links.__len__() > 0:
                                if not emission.outputs[0].links[0].to_node == materialout:
                                    make_links.append((emission.outputs[0], materialout.inputs[0]))
                            else:
                                make_links.append((emission.outputs[0], materialout.inputs[0]))
    
    
                            # Set brightness of viewer to compensate for Film and CM exposure
    
                            if context.scene.render.engine == 'CYCLES' and hasattr(context.scene, 'cycles'):
                                intensity = 1/context.scene.cycles.film_exposure  # Film exposure is a multiplier
    
                            intensity /= pow(2, (context.scene.view_settings.exposure))  # CM exposure is measured in stops/EVs (2^x)
                            emission.inputs[1].default_value = intensity
    
    
                            # Output type is 'SHADER', no Viewer needed. Delete Viewer if exists.
    
                            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 node in base_node_tree.nodes:
    
                                if node.name == 'Emission Viewer':
    
                                    delete_nodes.append((base_node_tree, node))
    
                        for li_from, li_to in make_links:
    
                            base_node_tree.links.new(li_from, li_to)
    
                        # Crate 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)
                            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:
                        if not self.is_socket_used_other_mats(socket):
                            tree = socket.id_data
                            tree.outputs.remove(socket)
    
                    # Delete nodes
    
                    for tree, node in delete_nodes:
                        tree.nodes.remove(node)
    
                    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=' '
        )
        color_prop: FloatVectorProperty(
            name="Color",
            description="The color of the frame node",
            default=(0.6, 0.6, 0.6),
            min=0, max=1, step=1, precision=3,
            subtype='COLOR_GAMMA', size=3
        )
    
    
        def execute(self, context):
            nodes, links = get_nodes_links(context)
            selected = []
            for node in nodes:
                if node.select == True:
                    selected.append(node)
    
            bpy.ops.node.add_node(type='NodeFrame')
            frm = nodes.active
            frm.label = self.label_prop
            frm.use_custom_color = True
            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'}
    
    
    Campbell Barton's avatar
    Campbell Barton committed
        to_type: EnumProperty(
    
            name="Switch to type",
            items=list(shaders_input_nodes_props) +
            list(shaders_output_nodes_props) +
            list(shaders_shader_nodes_props) +
            list(shaders_texture_nodes_props) +
            list(shaders_color_nodes_props) +
            list(shaders_vector_nodes_props) +
            list(shaders_converter_nodes_props) +
            list(shaders_layout_nodes_props) +
            list(compo_input_nodes_props) +
            list(compo_output_nodes_props) +
            list(compo_color_nodes_props) +
            list(compo_converter_nodes_props) +
            list(compo_filter_nodes_props) +
            list(compo_vector_nodes_props) +
            list(compo_matte_nodes_props) +
            list(compo_distort_nodes_props) +
    
            list(compo_layout_nodes_props) +
            list(blender_mat_input_nodes_props) +
            list(blender_mat_output_nodes_props) +
            list(blender_mat_color_nodes_props) +
            list(blender_mat_vector_nodes_props) +
            list(blender_mat_converter_nodes_props) +
            list(blender_mat_layout_nodes_props) +
            list(texture_input_nodes_props) +
            list(texture_output_nodes_props) +
            list(texture_color_nodes_props) +
            list(texture_pattern_nodes_props) +
            list(texture_textures_nodes_props) +
            list(texture_converter_nodes_props) +
            list(texture_distort_nodes_props) +
            list(texture_layout_nodes_props)
    
        geo_to_type: StringProperty(
            name="Switch to type",
            default = '',
        )
    
    
        def execute(self, context):
            nodes, links = get_nodes_links(context)
            to_type = self.to_type
    
            if self.geo_to_type != '':
                to_type = self.geo_to_type
    
            # 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': []},
                }
    
                for inout, soctype in (
                        ('INPUTS', 'MAIN',),
                        ('INPUTS', 'SHADER',),
                        ('INPUTS', 'RGBA',),
                        ('INPUTS', 'VECTOR',),
                        ('INPUTS', 'VALUE',),
                        ('OUTPUTS', 'MAIN',),
                        ('OUTPUTS', 'SHADER',),
                        ('OUTPUTS', 'RGBA',),
                        ('OUTPUTS', 'VECTOR',),
                        ('OUTPUTS', 'VALUE',),
                ):
                    if src_sockets[inout][soctype] and dst_sockets[inout][soctype]:
                        if soctype == 'MAIN':
                            sc = src_sockets[inout][src_sockets[inout]['MAIN']]
                            dt = dst_sockets[inout][dst_sockets[inout]['MAIN']]
                        else:
                            sc = src_sockets[inout][soctype]
                            dt = dst_sockets[inout][soctype]
                        # start with 'dt' to determine number of possibilities.
                        for i, soc in enumerate(dt):
                            # if src main has enough entries - match them with dst main sockets by indexes.
                            if len(sc) > i:
                                matches[inout][soctype].append(((sc[i][1], sc[i][3]), (soc[1], soc[3])))
                            # add 'VALUE_NAME' criterion to inputs.
                            if inout == 'INPUTS' and soctype == 'VALUE':
                                for s in sc:
                                    if s[2] == soc[2]:  # if names match
                                        # append src (index, dval), dst (index, dval)
                                        matches['INPUTS']['VALUE_NAME'].append(((s[1], s[3]), (soc[1], soc[3])))
    
                # When src ['INPUTS']['MAIN'] is 'VECTOR' replace 'MAIN' with matches VECTOR if possible.
                # This creates better links when relinking textures.
                if src_sockets['INPUTS']['MAIN'] == 'VECTOR' and matches['INPUTS']['VECTOR']:
                    matches['INPUTS']['MAIN'] = matches['INPUTS']['VECTOR']
    
                # Pass default values and RELINK:
                for tp in ('MAIN', 'SHADER', 'RGBA', 'VECTOR', 'VALUE_NAME', 'VALUE'):
                    # INPUTS: Base on matches in proper order.
                    for (src_i, src_dval), (dst_i, dst_dval) in matches['INPUTS'][tp]:
                        # pass dvals
                        if src_dval and dst_dval and tp in {'RGBA', 'VALUE_NAME'}:
                            new_node.inputs[dst_i].default_value = src_dval
                        # Special case: switch to math
                        if node.type in {'MIX_RGB', 'ALPHAOVER', 'ZCOMBINE'} and\
                                new_node.type == 'MATH' and\
                                tp == 'MAIN':
                            new_dst_dval = max(src_dval[0], src_dval[1], src_dval[2])
                            new_node.inputs[dst_i].default_value = new_dst_dval
                            if node.type == 'MIX_RGB':
                                if node.blend_type in [o[0] for o in operations]:
                                    new_node.operation = node.blend_type
                        # Special case: switch from math to some types
                        if node.type == 'MATH' and\
                                new_node.type in {'MIX_RGB', 'ALPHAOVER', 'ZCOMBINE'} and\
                                tp == 'MAIN':
                            for i in range(3):
                                new_node.inputs[dst_i].default_value[i] = src_dval
                            if new_node.type == 'MIX_RGB':
                                if node.operation in [t[0] for t in blend_types]:
                                    new_node.blend_type = node.operation
                                # Set Fac of MIX_RGB to 1.0
                                new_node.inputs[0].default_value = 1.0
                        # make link only when dst matching input is not linked already.
                        if node.inputs[src_i].links and not new_node.inputs[dst_i].links:
                            in_src_link = node.inputs[src_i].links[0]
                            in_dst_socket = new_node.inputs[dst_i]
                            links.new(in_src_link.from_socket, in_dst_socket)
                            links.remove(in_src_link)
                    # OUTPUTS: Base on matches in proper order.
                    for (src_i, src_dval), (dst_i, dst_dval) in matches['OUTPUTS'][tp]:
                        for out_src_link in node.outputs[src_i].links:
                            out_dst_socket = new_node.outputs[dst_i]
                            links.new(out_dst_socket, out_src_link.to_socket)
                # relink rest inputs if possible, no criteria
                for src_inp in node.inputs:
                    for dst_inp in new_node.inputs:
                        if src_inp.links and not dst_inp.links:
                            src_link = src_inp.links[0]
                            links.new(src_link.from_socket, dst_inp)
                            links.remove(src_link)
                # relink rest outputs if possible, base on node kind if any left.
                for src_o in node.outputs:
                    for out_src_link in src_o.links:
                        for dst_o in new_node.outputs:
                            if src_o.type == dst_o.type:
                                links.new(dst_o, out_src_link.to_socket)
                # relink rest outputs no criteria if any left. Link all from first output.
                for src_o in node.outputs:
                    for out_src_link in src_o.links:
                        if new_node.outputs:
                            links.new(new_node.outputs[0], out_src_link.to_socket)
                nodes.remove(node)
    
            force_update(context)
    
            return {'FINISHED'}
    
    
    class NWMergeNodes(Operator, NWBase):
        bl_idname = "node.nw_merge_nodes"
    
        bl_label = "Merge Nodes"
        bl_description = "Merge Selected Nodes"
        bl_options = {'REGISTER', 'UNDO'}
    
    
    Campbell Barton's avatar
    Campbell Barton committed
        mode: EnumProperty(
    
            description="All possible blend types, boolean operations and math operations",
            items= blend_types + [op for op in geo_combine_operations if op not in blend_types] + [op for op in operations if op not in blend_types],
    
    Campbell Barton's avatar
    Campbell Barton committed
        merge_type: EnumProperty(
    
            name="merge type",
            description="Type of Merge to be used",
            items=(
                ('AUTO', 'Auto', 'Automatic Output Type Detection'),
                ('SHADER', 'Shader', 'Merge using ADD or MIX Shader'),
    
                ('GEOMETRY', 'Geometry', 'Merge using Boolean or Join Geometry Node'),
    
                ('MIX', 'Mix Node', 'Merge using Mix Nodes'),
                ('MATH', 'Math Node', 'Merge using Math Nodes'),
    
                ('ZCOMBINE', 'Z-Combine Node', 'Merge using Z-Combine Nodes'),
                ('ALPHAOVER', 'Alpha Over Node', 'Merge using Alpha Over Nodes'),
    
        # Check if the link connects to a node that is in selected_nodes
        # If not, then check recursively for each link in the nodes outputs.
        # If yes, return True. If the recursion stops without finding a node
        # in selected_nodes, it returns False. The depth is used to prevent 
        # getting stuck in a loop because of an already present cycle.
        @staticmethod
        def link_creates_cycle(link, selected_nodes, depth=0)->bool:
            if depth > 255:
                # We're stuck in a cycle, but that cycle was already present,
                # so we return False. 
                # NOTE: The number 255 is arbitrary, but seems to work well.
                return False
            node = link.to_node
            if node in selected_nodes:
                return True
            if not node.outputs:
                return False
            for output in node.outputs:
                if output.is_linked:
                    for olink in output.links:
                        if NWMergeNodes.link_creates_cycle(olink, selected_nodes, depth+1):
                            return True
            # None of the outputs found a node in selected_nodes, so there is no cycle.
            return False
        
        # Merge the nodes in `nodes_list` with a node of type `node_name` that has a multi_input socket.
        # The parameters `socket_indices` gives the indices of the node sockets in the order that they should
        # be connected. The last one is assumed to be a multi input socket.
        # For convenience the node is returned.
        @staticmethod
        def merge_with_multi_input(nodes_list, merge_position,do_hide, loc_x, links, nodes, node_name, socket_indices):
            # The y-location of the last node
            loc_y = nodes_list[-1][2]
            if merge_position == 'CENTER':
                # Average the y-location
                for i in range(len(nodes_list)-1):
                    loc_y += nodes_list[i][2]
                loc_y = loc_y/len(nodes_list)
            new_node = nodes.new(node_name)
            new_node.hide = do_hide
            new_node.location.x = loc_x
            new_node.location.y = loc_y
            selected_nodes = [nodes[node_info[0]] for node_info in nodes_list]
            prev_links = []
            outputs_for_multi_input = []
            for i,node in enumerate(selected_nodes):
                node.select = False
                # Search for the first node which had output links that do not create
                # a cycle, which we can then reconnect afterwards.
                if prev_links == [] and node.outputs[0].is_linked:
                    prev_links = [link for link in node.outputs[0].links if not NWMergeNodes.link_creates_cycle(link, selected_nodes)]
                # Get the index of the socket, the last one is a multi input, and is thus used repeatedly
                # To get the placement to look right we need to reverse the order in which we connect the
                # outputs to the multi input socket.
                if i < len(socket_indices) - 1:
                    ind = socket_indices[i]
                    links.new(node.outputs[0], new_node.inputs[ind])
                else:
                    outputs_for_multi_input.insert(0, node.outputs[0])
            if outputs_for_multi_input != []:
                ind = socket_indices[-1]
                for output in outputs_for_multi_input:
                    links.new(output, new_node.inputs[ind])
            if prev_links != []:
                for link in prev_links:
                    links.new(new_node.outputs[0], link.to_node.inputs[0])
            return new_node
    
    
        def execute(self, context):
    
            settings = context.preferences.addons[__name__].preferences
    
            merge_hide = settings.merge_hide
            merge_position = settings.merge_position  # 'center' or 'bottom'
    
            do_hide = False
            do_hide_shader = False
            if merge_hide == 'ALWAYS':
                do_hide = True
                do_hide_shader = True
            elif merge_hide == 'NON_SHADER':
                do_hide = True
    
    
            tree_type = context.space_data.node_tree.type
    
            if tree_type == 'GEOMETRY':
                node_type = 'GeometryNode'
    
            if tree_type == 'COMPOSITING':
                node_type = 'CompositorNode'
            elif tree_type == 'SHADER':
                node_type = 'ShaderNode'
    
            elif tree_type == 'TEXTURE':
                node_type = 'TextureNode'
    
            nodes, links = get_nodes_links(context)
            mode = self.mode
            merge_type = self.merge_type
    
            # Prevent trying to add Z-Combine in not 'COMPOSITING' node tree.
            # 'ZCOMBINE' works only if mode == 'MIX'
            # Setting mode to None prevents trying to add 'ZCOMBINE' node.
    
            if (merge_type == 'ZCOMBINE' or merge_type == 'ALPHAOVER') and tree_type != 'COMPOSITING':
                merge_type = 'MIX'
                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':
                node_type = 'ShaderNode'
    
            selected_mix = []  # entry = [index, loc]
            selected_shader = []  # entry = [index, loc]
    
            selected_geometry = [] # entry = [index, loc]
    
            selected_math = []  # entry = [index, loc]
    
            selected_vector = [] # entry = [index, loc]
    
            selected_z = []  # entry = [index, loc]
    
            selected_alphaover = []  # entry = [index, loc]
    
    
            for i, node in enumerate(nodes):
                if node.select and node.outputs:
                    if merge_type == 'AUTO':
                        for (type, types_list, dst) in (
    
                                ('SHADER', ('MIX', 'ADD'), selected_shader),
    
                                ('GEOMETRY', [t[0] for t in geo_combine_operations], selected_geometry),
    
                                ('RGBA', [t[0] for t in blend_types], selected_mix),
                                ('VALUE', [t[0] for t in operations], selected_math),
    
                                ('VECTOR', [], selected_vector),
    
                            output_type = node.outputs[0].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.
                            if tree_type == 'GEOMETRY':
                                if mode == 'MIX':
                                    if output_type == 'VALUE' and type == 'VALUE':
                                        valid_mode = True
                                    elif output_type == 'VECTOR' and type == 'VECTOR':
                                        valid_mode = True
                                    elif type == 'GEOMETRY':
                                        valid_mode = True
    
                            # When mode is 'MIX' use mix node for both 'RGBA' and 'VALUE' output types.
                            # Cheat that output type is 'RGBA',
                            # and that 'MIX' exists in math operations list.
                            # This way when selected_mix list is analyzed:
                            # Node data will be appended even though it doesn't meet requirements.
    
                            elif output_type != 'SHADER' and mode == 'MIX':
    
                                output_type = 'RGBA'
                                valid_mode = True
                            if output_type == type and valid_mode:
    
                                dst.append([i, node.location.x, node.location.y, node.dimensions.x, node.hide])
    
                    else:
                        for (type, types_list, dst) in (
    
                                ('SHADER', ('MIX', 'ADD'), selected_shader),
    
                                ('GEOMETRY', [t[0] for t in geo_combine_operations], selected_geometry),
    
                                ('MIX', [t[0] for t in blend_types], selected_mix),
                                ('MATH', [t[0] for t in operations], selected_math),
    
                                ('ZCOMBINE', ('MIX', ), selected_z),
    
                                ('ALPHAOVER', ('MIX', ), selected_alphaover),
    
                            if merge_type == type and mode in types_list:
    
                                dst.append([i, node.location.x, node.location.y, node.dimensions.x, node.hide])
    
            # When nodes with output kinds 'RGBA' and 'VALUE' are selected at the same time
            # use only 'Mix' nodes for merging.
            # For that we add selected_math list to selected_mix list and clear selected_math.
            if selected_mix and selected_math and merge_type == 'AUTO':
                selected_mix += selected_math
                selected_math = []
    
            for nodes_list in [selected_mix, selected_shader, selected_geometry, selected_math, selected_vector, selected_z, selected_alphaover]:
                if not nodes_list:
                    continue
                count_before = len(nodes)
                # sort list by loc_x - reversed
                nodes_list.sort(key=lambda k: k[1], reverse=True)
                # get maximum loc_x
                loc_x = nodes_list[0][1] + nodes_list[0][3] + 70
                nodes_list.sort(key=lambda k: k[2], reverse=True)
                
                # 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:
                        node_type = 'ShaderNode'
                        if mode == 'MIX':
                            mode = 'ADD'
    
                        node_type = 'GeometryNode'
                if merge_position == 'CENTER':
                    loc_y = ((nodes_list[len(nodes_list) - 1][2]) + (nodes_list[len(nodes_list) - 2][2])) / 2  # average yloc of last two nodes (lowest two)
                    if nodes_list[len(nodes_list) - 1][-1] == True:  # if last node is hidden, mix should be shifted up a bit
                        if do_hide:
                            loc_y += 40
                        else:
                            loc_y += 80
                else:
                    loc_y = nodes_list[len(nodes_list) - 1][2]
                offset_y = 100
                if not do_hide:
                    offset_y = 200
                if nodes_list == selected_shader and not do_hide_shader:
                    offset_y = 150.0
                the_range = len(nodes_list) - 1
                if len(nodes_list) == 1:
                    the_range = 1
                was_multi = False
                for i in range(the_range):
                    if nodes_list == selected_mix:
                        add_type = node_type + 'MixRGB'
                        add = nodes.new(add_type)
                        add.blend_type = mode
                        if mode != 'MIX':
                            add.inputs[0].default_value = 1.0
                        add.show_preview = False
                        add.hide = do_hide
                        if do_hide:
                            loc_y = loc_y - 50
                        first = 1
                        second = 2
                        add.width_hidden = 100.0
                    elif nodes_list == selected_math:
                        add_type = node_type + 'Math'
                        add = nodes.new(add_type)
                        add.operation = mode
                        add.hide = do_hide
                        if do_hide:
                            loc_y = loc_y - 50
                        first = 0
                        second = 1
                        add.width_hidden = 100.0
                    elif nodes_list == selected_shader:
                        if mode == 'MIX':
                            add_type = node_type + 'MixShader'
    
                            add = nodes.new(add_type)
    
                            add.hide = do_hide_shader
                            if do_hide_shader:
    
                            first = 1
                            second = 2
                            add.width_hidden = 100.0
    
                        elif mode == 'ADD':
                            add_type = node_type + 'AddShader'
    
                            add = nodes.new(add_type)
    
                            add.hide = do_hide_shader
                            if do_hide_shader:
    
                            first = 0
                            second = 1
                            add.width_hidden = 100.0
    
                    elif nodes_list == selected_geometry:
                        if mode in ('JOIN', 'MIX'):
                            add_type = node_type + 'JoinGeometry'
                            add = self.merge_with_multi_input(nodes_list, merge_position, do_hide, loc_x, links, nodes, add_type,[0])
                        else:
                            add_type = node_type + 'Boolean'
                            indices = [0,1] if mode == 'DIFFERENCE' else [1]
                            add = self.merge_with_multi_input(nodes_list, merge_position, do_hide, loc_x, links, nodes, add_type,indices)
                            add.operation = mode
                        was_multi = True
                        break
                    elif nodes_list == selected_vector:
                        add_type = node_type + 'VectorMath'
                        add = nodes.new(add_type)
                        add.operation = mode
                        add.hide = do_hide
                        if do_hide:
                            loc_y = loc_y - 50
                        first = 0
                        second = 1
                        add.width_hidden = 100.0
                    elif nodes_list == selected_z:
                        add = nodes.new('CompositorNodeZcombine')
                        add.show_preview = False
                        add.hide = do_hide
                        if do_hide:
                            loc_y = loc_y - 50
                        first = 0
                        second = 2
                        add.width_hidden = 100.0
                    elif nodes_list == selected_alphaover:
                        add = nodes.new('CompositorNodeAlphaOver')
                        add.show_preview = False
                        add.hide = do_hide
                        if do_hide:
                            loc_y = loc_y - 50
                        first = 1
                        second = 2
                        add.width_hidden = 100.0
                    add.location = loc_x, loc_y
                    loc_y += offset_y
                    add.select = True
                
                # This has already been handled separately
                if was_multi:
                    continue
                count_adds = i + 1
                count_after = len(nodes)
                index = count_after - 1
                first_selected = nodes[nodes_list[0][0]]
                # "last" node has been added as first, so its index is count_before.
                last_add = nodes[count_before]
                # Create list of invalid indexes.
                invalid_nodes = [nodes[n[0]] for n in (selected_mix + selected_math + selected_shader + selected_z + selected_geometry)]
                
                # 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.
                if len(nodes_list) == 2:
                    if not first_selected.outputs[0].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.
                            # Link only if "to_node" index not in invalid indexes list.
                            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:
                    # 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])
                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'):
                            links.new(fs_out, node_to.inputs[1])
                            break
                # add links between added ADD nodes and between selected and ADD nodes
                for i in range(count_adds):
                    if i < count_adds - 1:
                        node_from = nodes[index]
                        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])
                        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'):
                                    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])
                        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'):
                                    links.new(from_out, node_to.inputs[node_to_z_i])
                    index -= 1
                # set "last" of added nodes as active
                nodes.active = last_add
                for i, x, y, dx, h in nodes_list:
                    nodes[i].select = False
    
    class NWBatchChangeNodes(Operator, NWBase):
        bl_idname = "node.nw_batch_change"
    
        bl_label = "Batch Change"
        bl_description = "Batch Change Blend Type and Math Operation"
        bl_options = {'REGISTER', 'UNDO'}
    
    
    Campbell Barton's avatar
    Campbell Barton committed
        blend_type: EnumProperty(
    
            name="Blend Type",
            items=blend_types + navs,
        )
    
    Campbell Barton's avatar
    Campbell Barton committed
        operation: EnumProperty(
    
            name="Operation",
            items=operations + navs,
        )
    
    
        def execute(self, context):
            blend_type = self.blend_type
            operation = self.operation
            for node in context.selected_nodes:
    
                if node.type == 'MIX_RGB' or node.bl_idname == 'GeometryNodeAttributeMix':
    
                    if not blend_type in [nav[0] for nav in navs]:
                        node.blend_type = blend_type
                    else:
                        if blend_type == 'NEXT':
                            index = [i for i, entry in enumerate(blend_types) if node.blend_type in entry][0]
                            #index = blend_types.index(node.blend_type)
                            if index == len(blend_types) - 1:
                                node.blend_type = blend_types[0][0]
                            else:
                                node.blend_type = blend_types[index + 1][0]
    
                        if blend_type == 'PREV':
                            index = [i for i, entry in enumerate(blend_types) if node.blend_type in entry][0]
                            if index == 0:
                                node.blend_type = blend_types[len(blend_types) - 1][0]
                            else:
                                node.blend_type = blend_types[index - 1][0]
    
    
                if node.type == 'MATH' or node.bl_idname == 'GeometryNodeAttributeMath':
    
                    if not operation in [nav[0] for nav in navs]:
                        node.operation = operation
                    else:
                        if operation == 'NEXT':
                            index = [i for i, entry in enumerate(operations) if node.operation in entry][0]
                            #index = operations.index(node.operation)
                            if index == len(operations) - 1:
                                node.operation = operations[0][0]
                            else:
                                node.operation = operations[index + 1][0]
    
                        if operation == 'PREV':
                            index = [i for i, entry in enumerate(operations) if node.operation in entry][0]
                            #index = operations.index(node.operation)
                            if index == 0:
                                node.operation = operations[len(operations) - 1][0]
                            else:
                                node.operation = operations[index - 1][0]
    
            return {'FINISHED'}
    
    
    
    class NWChangeMixFactor(Operator, NWBase):
        bl_idname = "node.nw_factor"
    
        bl_label = "Change Factor"
        bl_description = "Change Factors of Mix Nodes and Mix Shader Nodes"
        bl_options = {'REGISTER', 'UNDO'}
    
        # option: Change factor.
        # If option is 1.0 or 0.0 - set to 1.0 or 0.0
        # Else - change factor by option value.
    
    Campbell Barton's avatar
    Campbell Barton committed
        option: FloatProperty()
    
    
        def execute(self, context):
            nodes, links = get_nodes_links(context)
            option = self.option
            selected = []  # entry = index
            for si, node in enumerate(nodes):
                if node.select:
                    if node.type in {'MIX_RGB', 'MIX_SHADER'}:
                        selected.append(si)
    
            for si in selected:
                fac = nodes[si].inputs[0]
                nodes[si].hide = False
                if option in {0.0, 1.0}:
                    fac.default_value = option
                else:
                    fac.default_value += option
    
            return {'FINISHED'}
    
    
    
    class NWCopySettings(Operator, NWBase):
        bl_idname = "node.nw_copy_settings"
    
        bl_label = "Copy Settings"
        bl_description = "Copy Settings of Active Node to Selected Nodes"
        bl_options = {'REGISTER', 'UNDO'}
    
        @classmethod
        def poll(cls, context):
            valid = False
    
            if nw_check(context):
    
                if (
                        context.active_node is not None and
                        context.active_node.type != 'FRAME'
                ):
    
            return valid
    
        def execute(self, context):
    
            node_active = context.active_node
            node_selected = context.selected_nodes
    
            # Error handling
            if not (len(node_selected) > 1):
                self.report({'ERROR'}, "2 nodes must be selected at least")
                return {'CANCELLED'}
    
            # Check if active node is in the selection
            selected_node_names = [n.name for n in node_selected]
            if node_active.name not in selected_node_names:
                self.report({'ERROR'}, "No active node")
                return {'CANCELLED'}
    
            # Get nodes in selection by type
            valid_nodes = [n for n in node_selected if n.type == node_active.type]
    
            if not (len(valid_nodes) > 1) and node_active:
                self.report({'ERROR'}, "Selected nodes are not of the same type as {}".format(node_active.name))
                return {'CANCELLED'}
    
            if len(valid_nodes) != len(node_selected):
                # Report nodes that are not valid
                valid_node_names = [n.name for n in valid_nodes]
                not_valid_names = list(set(selected_node_names) - set(valid_node_names))
                self.report({'INFO'}, "Ignored {} (not of the same type as {})".format(", ".join(not_valid_names), node_active.name))
    
    
            # Reference original
    
            orig = node_active
            #node_selected_names = [n.name for n in node_selected]
    
            # Output list
            success_names = []
    
            # Deselect all nodes
            for i in node_selected:
                i.select = False
    
            # Code by zeffii from http://blender.stackexchange.com/a/42338/3710
    
            # Run through all other nodes
            for node in valid_nodes[1:]:
    
                # Check for frame node
                parent = node.parent if node.parent else None
                node_loc = [node.location.x, node.location.y]
    
                # Select original to duplicate
                orig.select = True
    
                # Duplicate selected node
                bpy.ops.node.duplicate()
                new_node = context.selected_nodes[0]
    
                # Deselect copy
    
                new_node.select = False
    
    
                # Properties to copy
                node_tree = node.id_data
                props_to_copy = 'bl_idname name location height width'.split(' ')
    
                # Input and outputs
                reconnections = []
                mappings = chain.from_iterable([node.inputs, node.outputs])
                for i in (i for i in mappings if i.is_linked):
                    for L in i.links:
                        reconnections.append([L.from_socket.path_from_id(), L.to_socket.path_from_id()])
    
                # Properties
                props = {j: getattr(node, j) for j in props_to_copy}
                props_to_copy.pop(0)
    
                for prop in props_to_copy:
                    setattr(new_node, prop, props[prop])
    
                # Get the node tree to remove the old node
                nodes = node_tree.nodes
                nodes.remove(node)
                new_node.name = props['name']
    
                if parent:
                    new_node.parent = parent
                    new_node.location = node_loc
    
                for str_from, str_to in reconnections:
                    node_tree.links.new(eval(str_from), eval(str_to))
    
                success_names.append(new_node.name)
    
            orig.select = True
            node_tree.nodes.active = orig
            self.report({'INFO'}, "Successfully copied attributes from {} to: {}".format(orig.name, ", ".join(success_names)))
    
    class NWCopyLabel(Operator, NWBase):
        bl_idname = "node.nw_copy_label"
    
        bl_label = "Copy Label"
        bl_options = {'REGISTER', 'UNDO'}
    
    
    Campbell Barton's avatar
    Campbell Barton committed
        option: EnumProperty(
    
            name="option",
            description="Source of name of label",
            items=(
                ('FROM_ACTIVE', 'from active', 'from active node',),
                ('FROM_NODE', 'from node', 'from node linked to selected node'),
                ('FROM_SOCKET', 'from socket', 'from socket linked to selected node'),
            )
        )
    
    
        def execute(self, context):
            nodes, links = get_nodes_links(context)