Skip to content
Snippets Groups Projects
node_wrangler.py 187 KiB
Newer Older
  • Learn to ignore specific revisions
  •         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 MixRGB node and math nodes used for geometry nodes are of type 'ShaderNode'
            if (merge_type == 'MATH' or merge_type == 'MIX') 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 = get_first_enabled_output(node)
                            output_type = output.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 or nodes_list is selected_mix:
    
                        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.
    
                first_selected_output = get_first_enabled_output(first_selected)
    
                if len(nodes_list) == 2:
    
                    if not first_selected_output.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_output.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_output, node_to.inputs[first])
    
                if node_to.type == 'ZCOMBINE':
                    for fs_out in first_selected.outputs:
    
                        if fs_out != first_selected_output 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(get_first_enabled_output(node_from), node_to.inputs[node_to_input_i])
    
                        if node_to.type == 'ZCOMBINE':
                            for from_out in node_from.outputs:
    
                                if from_out != get_first_enabled_output(node_from) 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(get_first_enabled_output(node_from), node_to.inputs[node_to_input_i])
    
                        if node_to.type == 'ZCOMBINE':
                            for from_out in node_from.outputs:
    
                                if from_out != get_first_enabled_output(node_from) 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)
            option = self.option
            active = nodes.active
            if option == 'FROM_ACTIVE':
                if active:
                    src_label = active.label
                    for node in [n for n in nodes if n.select and nodes.active != n]:
                        node.label = src_label
            elif option == 'FROM_NODE':
                selected = [n for n in nodes if n.select]
                for node in selected:
                    for input in node.inputs:
                        if input.links:
                            src = input.links[0].from_node
                            node.label = src.label
                            break
            elif option == 'FROM_SOCKET':
                selected = [n for n in nodes if n.select]
                for node in selected:
                    for input in node.inputs:
                        if input.links:
                            src = input.links[0].from_socket
                            node.label = src.name
                            break
    
            return {'FINISHED'}
    
    
    
    class NWClearLabel(Operator, NWBase):
        bl_idname = "node.nw_clear_label"
    
        bl_label = "Clear Label"
        bl_options = {'REGISTER', 'UNDO'}
    
    
    Campbell Barton's avatar
    Campbell Barton committed
        option: BoolProperty()
    
    
        def execute(self, context):
            nodes, links = get_nodes_links(context)
            for node in [n for n in nodes if n.select]:
                node.label = ''
    
            return {'FINISHED'}
    
        def invoke(self, context, event):
            if self.option:
                return self.execute(context)
            else:
                return context.window_manager.invoke_confirm(self, event)
    
    
    
    class NWModifyLabels(Operator, NWBase):
    
        """Modify Labels of all selected nodes"""
    
        bl_idname = "node.nw_modify_labels"
        bl_label = "Modify Labels"
        bl_options = {'REGISTER', 'UNDO'}
    
    
        replace_from: StringProperty(
    
        replace_to: StringProperty(
    
            name="Replace with"
        )
    
        def execute(self, context):
            nodes, links = get_nodes_links(context)
            for node in [n for n in nodes if n.select]:
                node.label = self.prepend + node.label.replace(self.replace_from, self.replace_to) + self.append
    
            return {'FINISHED'}
    
        def invoke(self, context, event):
            self.prepend = ""
            self.append = ""
            self.remove = ""
            return context.window_manager.invoke_props_dialog(self)
    
    
    class NWAddTextureSetup(Operator, NWBase):
        bl_idname = "node.nw_add_texture"
    
        bl_label = "Texture Setup"
        bl_description = "Add Texture Node Setup to Selected Shaders"
        bl_options = {'REGISTER', 'UNDO'}
    
    
    Campbell Barton's avatar
    Campbell Barton committed
        add_mapping: BoolProperty(name="Add Mapping Nodes", description="Create coordinate and mapping nodes for the texture (ignored for selected texture nodes)", default=True)
    
        @classmethod
        def poll(cls, context):
    
            if nw_check(context):
                space = context.space_data
    
    
        def execute(self, context):
            nodes, links = get_nodes_links(context)
    
    
            texture_types = [x.nodetype for x in
                             get_nodes_from_category('Texture', context)]
    
            selected_nodes = [n for n in nodes if n.select]
    
            for node in selected_nodes:
                if not node.inputs:
                    continue
    
    
                target_input = node.inputs[0]
                for input in node.inputs:
                    if input.enabled:
                        input_index += 1
                        if not input.is_linked:
                            target_input = input
    
                    self.report({'WARNING'}, "No free inputs for node: " + node.name)
                    continue
    
                x_offset = 0
                padding = 40.0
                locx = node.location.x
                locy = node.location.y - (input_index * padding)
    
    
                is_texture_node = node.rna_type.identifier in texture_types
    
                use_environment_texture = node.type == 'BACKGROUND'
    
                # Add an image texture before normal shader nodes.
                if not is_texture_node:
                    image_texture_type = 'ShaderNodeTexEnvironment' if use_environment_texture else 'ShaderNodeTexImage'
                    image_texture_node = nodes.new(image_texture_type)
                    x_offset = x_offset + image_texture_node.width + padding
                    image_texture_node.location = [locx - x_offset, locy]
                    nodes.active = image_texture_node
                    links.new(image_texture_node.outputs[0], target_input)
    
                    # The mapping setup following this will connect to the firrst input of this image texture.
                    target_input = image_texture_node.inputs[0]
    
                node.select = False
    
                if is_texture_node or self.add_mapping:
                    # Add Mapping node.
                    mapping_node = nodes.new('ShaderNodeMapping')
                    x_offset = x_offset + mapping_node.width + padding
                    mapping_node.location = [locx - x_offset, locy]
                    links.new(mapping_node.outputs[0], target_input)
    
                    # Add Texture Coordinates node.
                    tex_coord_node = nodes.new('ShaderNodeTexCoord')
                    x_offset = x_offset + tex_coord_node.width + padding
                    tex_coord_node.location = [locx - x_offset, locy]
    
                    is_procedural_texture = is_texture_node and node.type != 'TEX_IMAGE'
                    use_generated_coordinates = is_procedural_texture or use_environment_texture
                    tex_coord_output = tex_coord_node.outputs[0 if use_generated_coordinates else 2]
                    links.new(tex_coord_output, mapping_node.inputs[0])
    
    
    class NWAddPrincipledSetup(Operator, NWBase, ImportHelper):
        bl_idname = "node.nw_add_textures_for_principled"
        bl_label = "Principled Texture Setup"
        bl_description = "Add Texture Node Setup for Principled BSDF"
        bl_options = {'REGISTER', 'UNDO'}
    
    
        directory: StringProperty(
            name='Directory',
            subtype='DIR_PATH',
            default='',
            description='Folder to search in for image files'
        )
        files: CollectionProperty(
            type=bpy.types.OperatorFileListElement,
            options={'HIDDEN', 'SKIP_SAVE'}
        )
    
        relative_path: BoolProperty(
            name='Relative Path',
    
            description='Set the file path relative to the blend file, when possible',
    
        def draw(self, context):
            layout = self.layout
            layout.alignment = 'LEFT'
    
            layout.prop(self, 'relative_path')
    
    
        @classmethod
        def poll(cls, context):
            valid = False
            if nw_check(context):
                space = context.space_data
    
                    valid = True
            return valid
    
        def execute(self, context):
            # Check if everything is ok
            if not self.directory:
                self.report({'INFO'}, 'No Folder Selected')
                return {'CANCELLED'}
            if not self.files[:]:
                self.report({'INFO'}, 'No Files Selected')
                return {'CANCELLED'}
    
            nodes, links = get_nodes_links(context)
            active_node = nodes.active
    
            if not (active_node and active_node.bl_idname == 'ShaderNodeBsdfPrincipled'):
    
                self.report({'INFO'}, 'Select Principled BSDF')
                return {'CANCELLED'}
    
            # Helper_functions
            def split_into__components(fname):
                # Split filename into components
                # 'WallTexture_diff_2k.002.jpg' -> ['Wall', 'Texture', 'diff', 'k']
                # Remove extension
                fname = path.splitext(fname)[0]
                # Remove digits
                fname = ''.join(i for i in fname if not i.isdigit())
    
                fname = re.sub(r"([a-z])([A-Z])", r"\g<1> \g<2>",fname)
    
                # Replace common separators with SPACE
    
                separators = ['_', '.', '-', '__', '--', '#']
                for sep in separators:
    
                    fname = fname.replace(sep, ' ')
    
                components = fname.split(' ')
                components = [c.lower() for c in components]
                return components
    
            # Filter textures names for texturetypes in filenames
            # [Socket Name, [abbreviations and keyword list], Filename placeholder]
    
            tags = context.preferences.addons[__name__].preferences.principled_tags
    
            normal_abbr = tags.normal.split(' ')
            bump_abbr = tags.bump.split(' ')
            gloss_abbr = tags.gloss.split(' ')
            rough_abbr = tags.rough.split(' ')
    
            ['Displacement', tags.displacement.split(' '), None],
            ['Base Color', tags.base_color.split(' '), None],
            ['Subsurface Color', tags.sss_color.split(' '), None],
            ['Metallic', tags.metallic.split(' '), None],
            ['Specular', tags.specular.split(' '), None],
    
            ['Roughness', rough_abbr + gloss_abbr, None],
            ['Normal', normal_abbr + bump_abbr, None],
    
            ['Transmission', tags.transmission.split(' '), None],
            ['Emission', tags.emission.split(' '), None],
            ['Alpha', tags.alpha.split(' '), None],
            ['Ambient Occlusion', tags.ambient_occlusion.split(' '), None],
    
            # Look through texture_types and set value as filename of first matched file
            def match_files_to_socket_names():
                for sname in socketnames:
                    for file in self.files:
                        fname = file.name
                        filenamecomponents = split_into__components(fname)
                        matches = set(sname[1]).intersection(set(filenamecomponents))
    
                        # TODO: ignore basename (if texture is named "fancy_metal_nor", it will be detected as metallic map, not normal map)
    
                        if matches:
                            sname[2] = fname
                            break
    
            match_files_to_socket_names()
            # Remove socketnames without found files
            socketnames = [s for s in socketnames if s[2]
                           and path.exists(self.directory+s[2])]
            if not socketnames:
                self.report({'INFO'}, 'No matching images found')
                print('No matching images found')
                return {'CANCELLED'}
    
    
            # Don't override path earlier as os.path is used to check the absolute path
            import_path = self.directory
            if self.relative_path:
                if bpy.data.filepath:
    
                    try:
                        import_path = bpy.path.relpath(self.directory)
                    except ValueError:
                        pass
    
            # Add found images
            print('\nMatched Textures:')
            texture_nodes = []
            disp_texture = None
    
            normal_node = None
            roughness_node = None
            for i, sname in enumerate(socketnames):
                print(i, sname[0], sname[2])
    
                # DISPLACEMENT NODES
                if sname[0] == 'Displacement':
                    disp_texture = nodes.new(type='ShaderNodeTexImage')
    
                    img = bpy.data.images.load(path.join(import_path, sname[2]))
    
                    disp_texture.image = img
                    disp_texture.label = 'Displacement'
    
                    if disp_texture.image:
                        disp_texture.image.colorspace_settings.is_data = True
    
    
                    # Add displacement offset nodes
    
                    disp_node = nodes.new(type='ShaderNodeDisplacement')
    
                    # Align the Displacement node under the active Principled BSDF node
                    disp_node.location = active_node.location + Vector((100, -700))
    
                    link = links.new(disp_node.inputs[0], disp_texture.outputs[0])
    
                    # TODO Turn on true displacement in the material
    
                    # Find output node
    
                    output_node = [n for n in nodes if n.bl_idname == 'ShaderNodeOutputMaterial']
                    if output_node:
                        if not output_node[0].inputs[2].is_linked:
    
                            link = links.new(output_node[0].inputs[2], disp_node.outputs[0])
    
                # AMBIENT OCCLUSION TEXTURE
                if sname[0] == 'Ambient Occlusion':
                    ao_texture = nodes.new(type='ShaderNodeTexImage')
                    img = bpy.data.images.load(path.join(import_path, sname[2]))
                    ao_texture.image = img
                    ao_texture.label = sname[0]
                    if ao_texture.image:
                        ao_texture.image.colorspace_settings.is_data = True
    
                    continue
    
    
                if not active_node.inputs[sname[0]].is_linked:
                    # No texture node connected -> add texture node with new image
                    texture_node = nodes.new(type='ShaderNodeTexImage')
    
                    img = bpy.data.images.load(path.join(import_path, sname[2]))
    
                    texture_node.image = img
    
                    # NORMAL NODES
                    if sname[0] == 'Normal':
                        # Test if new texture node is normal or bump map
                        fname_components = split_into__components(sname[2])
                        match_normal = set(normal_abbr).intersection(set(fname_components))
                        match_bump = set(bump_abbr).intersection(set(fname_components))
                        if match_normal:
                            # If Normal add normal node in between
                            normal_node = nodes.new(type='ShaderNodeNormalMap')
                            link = links.new(normal_node.inputs[1], texture_node.outputs[0])
                        elif match_bump:
                            # If Bump add bump node in between
                            normal_node = nodes.new(type='ShaderNodeBump')
                            link = links.new(normal_node.inputs[2], texture_node.outputs[0])
    
                        link = links.new(active_node.inputs[sname[0]], normal_node.outputs[0])
                        normal_node_texture = texture_node
    
                    elif sname[0] == 'Roughness':
                        # Test if glossy or roughness map
                        fname_components = split_into__components(sname[2])
                        match_rough = set(rough_abbr).intersection(set(fname_components))
                        match_gloss = set(gloss_abbr).intersection(set(fname_components))
    
                        if match_rough:
                            # If Roughness nothing to to
                            link = links.new(active_node.inputs[sname[0]], texture_node.outputs[0])
    
                        elif match_gloss:
                            # If Gloss Map add invert node
                            invert_node = nodes.new(type='ShaderNodeInvert')
                            link = links.new(invert_node.inputs[1], texture_node.outputs[0])
    
                            link = links.new(active_node.inputs[sname[0]], invert_node.outputs[0])
                            roughness_node = texture_node
    
                    else:
                        # This is a simple connection Texture --> Input slot
                        link = links.new(active_node.inputs[sname[0]], texture_node.outputs[0])
    
                    # Use non-color for all but 'Base Color' Textures
    
                    if not sname[0] in ['Base Color', 'Emission'] and texture_node.image:
    
                        texture_node.image.colorspace_settings.is_data = True
    
    
                else:
                    # If already texture connected. add to node list for alignment
                    texture_node = active_node.inputs[sname[0]].links[0].from_node
    
                # This are all connected texture nodes
                texture_nodes.append(texture_node)
                texture_node.label = sname[0]
    
            if disp_texture:
                texture_nodes.append(disp_texture)
    
    
            if ao_texture:
                # We want the ambient occlusion texture to be the top most texture node
                texture_nodes.insert(0, ao_texture)
    
    
            # Alignment
            for i, texture_node in enumerate(texture_nodes):
    
                offset = Vector((-550, (i * -280) + 200))
    
                texture_node.location = active_node.location + offset
    
            if normal_node:
                # Extra alignment if normal node was added
    
                normal_node.location = normal_node_texture.location + Vector((300, 0))
    
    
            if roughness_node:
                # Alignment of invert node if glossy map
    
                invert_node.location = roughness_node.location + Vector((300, 0))
    
    
            # Add texture input + mapping
            mapping = nodes.new(type='ShaderNodeMapping')
    
            mapping.location = active_node.location + Vector((-1050, 0))
    
            if len(texture_nodes) > 1:
                # If more than one texture add reroute node in between
                reroute = nodes.new(type='NodeReroute')
    
                tex_coords = Vector((texture_nodes[0].location.x, sum(n.location.y for n in texture_nodes)/len(texture_nodes)))
                reroute.location = tex_coords + Vector((-50, -120))
                for texture_node in texture_nodes:
                    link = links.new(texture_node.inputs[0], reroute.outputs[0])
                link = links.new(reroute.inputs[0], mapping.outputs[0])
            else:
                link = links.new(texture_nodes[0].inputs[0], mapping.outputs[0])
    
            # Connect texture_coordiantes to mapping node
            texture_input = nodes.new(type='ShaderNodeTexCoord')
            texture_input.location = mapping.location + Vector((-200, 0))
            link = links.new(mapping.inputs[0], texture_input.outputs[2])
    
    
            # Create frame around tex coords and mapping
            frame = nodes.new(type='NodeFrame')
            frame.label = 'Mapping'
            mapping.parent = frame
            texture_input.parent = frame
            frame.update()
    
            # Create frame around texture nodes
            frame = nodes.new(type='NodeFrame')
            frame.label = 'Textures'
            for tnode in texture_nodes:
                tnode.parent = frame
            frame.update()
    
    
            # Just to be sure
            active_node.select = False
            nodes.update()
            links.update()
            force_update(context)
            return {'FINISHED'}
    
    class NWAddReroutes(Operator, NWBase):
        """Add Reroute Nodes and link them to outputs of selected nodes"""
        bl_idname = "node.nw_add_reroutes"
    
        bl_label = "Add Reroutes"
        bl_description = "Add Reroutes to Outputs"
        bl_options = {'REGISTER', 'UNDO'}
    
    
    Campbell Barton's avatar
    Campbell Barton committed
        option: EnumProperty(
    
            name="option",
            items=[
                ('ALL', 'to all', 'Add to all outputs'),
                ('LOOSE', 'to loose', 'Add only to loose outputs'),
                ('LINKED', 'to linked', 'Add only to linked outputs'),
            ]
        )
    
    
        def execute(self, context):
            tree_type = context.space_data.node_tree.type
            option = self.option
            nodes, links = get_nodes_links(context)
            # output valid when option is 'all' or when 'loose' output has no links
            valid = False
            post_select = []  # nodes to be selected after execution
            # create reroutes and recreate links
            for node in [n for n in nodes if n.select]:
                if node.outputs:
                    x = node.location.x
                    y = node.location.y
                    width = node.width
                    # unhide 'REROUTE' nodes to avoid issues with location.y
                    if node.type == 'REROUTE':
                        node.hide = False
                    # When node is hidden - width_hidden not usable.
                    # Hack needed to calculate real width
                    if node.hide:
                        bpy.ops.node.select_all(action='DESELECT')
                        helper = nodes.new('NodeReroute')
                        helper.select = True
                        node.select = True
                        # resize node and helper to zero. Then check locations to calculate width
                        bpy.ops.transform.resize(value=(0.0, 0.0, 0.0))
                        width = 2.0 * (helper.location.x - node.location.x)
                        # restore node location
                        node.location = x, y
                        # delete helper
                        node.select = False
                        # only helper is selected now
                        bpy.ops.node.delete()