Skip to content
Snippets Groups Projects
node_wrangler.py 187 KiB
Newer Older
  • Learn to ignore specific revisions
  •                 x = node.location.x + width + 20.0
                    if node.type != 'REROUTE':
                        y -= 35.0
                    y_offset = -22.0
                    loc = x, y
                reroutes_count = 0  # will be used when aligning reroutes added to hidden nodes
                for out_i, output in enumerate(node.outputs):
                    pass_used = False  # initial value to be analyzed if 'R_LAYERS'
    
                    # if node != 'R_LAYERS' - "pass_used" not needed, so set it to True
    
                    if node.type != 'R_LAYERS':
                        pass_used = True
                    else:  # if 'R_LAYERS' check if output represent used render pass
                        node_scene = node.scene
                        node_layer = node.layer
                        # If output - "Alpha" is analyzed - assume it's used. Not represented in passes.
                        if output.name == 'Alpha':
                            pass_used = True
                        else:
                            # check entries in global 'rl_outputs' variable
    
                                if output.name in {rlo.output_name, rlo.exr_output_name}:
                                    pass_used = getattr(node_scene.view_layers[node_layer], rlo.render_pass)
    
                                    break
                    if pass_used:
                        valid = ((option == 'ALL') or
                                 (option == 'LOOSE' and not output.links) or
                                 (option == 'LINKED' and output.links))
                        # Add reroutes only if valid, but offset location in all cases.
                        if valid:
                            n = nodes.new('NodeReroute')
                            nodes.active = n
                            for link in output.links:
                                links.new(n.outputs[0], link.to_socket)
                            links.new(output, n.inputs[0])
                            n.location = loc
                            post_select.append(n)
                        reroutes_count += 1
                        y += y_offset
                        loc = x, y
                # disselect the node so that after execution of script only newly created nodes are selected
                node.select = False
                # nicer reroutes distribution along y when node.hide
                if node.hide:
                    y_translate = reroutes_count * y_offset / 2.0 - y_offset - 35.0
                    for reroute in [r for r in nodes if r.select]:
                        reroute.location.y -= y_translate
                for node in post_select:
                    node.select = True
    
            return {'FINISHED'}
    
    
    
    class NWLinkActiveToSelected(Operator, NWBase):
        """Link active node to selected nodes basing on various criteria"""
        bl_idname = "node.nw_link_active_to_selected"
    
        bl_label = "Link Active Node to Selected"
        bl_options = {'REGISTER', 'UNDO'}
    
    
    Campbell Barton's avatar
    Campbell Barton committed
        replace: BoolProperty()
        use_node_name: BoolProperty()
        use_outputs_names: BoolProperty()
    
    
        @classmethod
        def poll(cls, context):
            valid = False
    
            if nw_check(context):
                if context.active_node is not None:
    
                    if context.active_node.select:
                        valid = True
            return valid
    
        def execute(self, context):
            nodes, links = get_nodes_links(context)
            replace = self.replace
            use_node_name = self.use_node_name
            use_outputs_names = self.use_outputs_names
            active = nodes.active
            selected = [node for node in nodes if node.select and node != active]
            outputs = []  # Only usable outputs of active nodes will be stored here.
            for out in active.outputs:
                if active.type != 'R_LAYERS':
                    outputs.append(out)
                else:
                    # 'R_LAYERS' node type needs special handling.
                    # outputs of 'R_LAYERS' are callable even if not seen in UI.
                    # Only outputs that represent used passes should be taken into account
                    # Check if pass represented by output is used.
                    # global 'rl_outputs' list will be used for that
    
                    for rlo in rl_outputs:
    
                        pass_used = False  # initial value. Will be set to True if pass is used
                        if out.name == 'Alpha':
                            # Alpha output is always present. Doesn't have representation in render pass. Assume it's used.
                            pass_used = True
    
                        elif out.name in {rlo.output_name, rlo.exr_output_name}:
    
                            # example 'render_pass' entry: 'use_pass_uv' Check if True in scene render layers
    
                            pass_used = getattr(active.scene.view_layers[active.layer], rlo.render_pass)
    
                            break
                    if pass_used:
                        outputs.append(out)
            doit = True  # Will be changed to False when links successfully added to previous output.
            for out in outputs:
                if doit:
                    for node in selected:
                        dst_name = node.name  # Will be compared with src_name if needed.
                        # When node has label - use it as dst_name
                        if node.label:
                            dst_name = node.label
                        valid = True  # Initial value. Will be changed to False if names don't match.
    
                        src_name = dst_name  # If names not used - this assignment will keep valid = True.
    
                        if use_node_name:
                            # Set src_name to source node name or label
                            src_name = active.name
                            if active.label:
                                src_name = active.label
                        elif use_outputs_names:
                            src_name = (out.name, )
    
                            for rlo in rl_outputs:
                                if out.name in {rlo.output_name, rlo.exr_output_name}:
                                    src_name = (rlo.output_name, rlo.exr_output_name)
    
                        if dst_name not in src_name:
                            valid = False
                        if valid:
                            for input in node.inputs:
                                if input.type == out.type or node.type == 'REROUTE':
                                    if replace or not input.is_linked:
                                        links.new(out, input)
                                        if not use_node_name and not use_outputs_names:
                                            doit = False
                                        break
    
            return {'FINISHED'}
    
    
    
    class NWAlignNodes(Operator, NWBase):
    
        '''Align the selected nodes neatly in a row/column'''
    
        bl_idname = "node.nw_align_nodes"
    
        bl_label = "Align Nodes"
    
        bl_options = {'REGISTER', 'UNDO'}
    
    Campbell Barton's avatar
    Campbell Barton committed
        margin: IntProperty(name='Margin', default=50, description='The amount of space between nodes')
    
    
        def execute(self, context):
            nodes, links = get_nodes_links(context)
    
            margin = self.margin
    
            selection = []
            for node in nodes:
                if node.select and node.type != 'FRAME':
                    selection.append(node)
    
            # If no nodes are selected, align all nodes
    
            if not selection:
                selection = nodes
    
            elif nodes.active in selection:
                active_loc = copy(nodes.active.location)  # make a copy, not a reference
    
            # Check if nodes should be laid out horizontally or vertically
    
            x_locs = [n.location.x + (n.dimensions.x / 2) for n in selection]  # use dimension to get center of node, not corner
            y_locs = [n.location.y - (n.dimensions.y / 2) for n in selection]
            x_range = max(x_locs) - min(x_locs)
            y_range = max(y_locs) - min(y_locs)
            mid_x = (max(x_locs) + min(x_locs)) / 2
            mid_y = (max(y_locs) + min(y_locs)) / 2
            horizontal = x_range > y_range
    
            # Sort selection by location of node mid-point
            if horizontal:
                selection = sorted(selection, key=lambda n: n.location.x + (n.dimensions.x / 2))
            else:
                selection = sorted(selection, key=lambda n: n.location.y - (n.dimensions.y / 2), reverse=True)
    
            # Alignment
            current_pos = 0
            for node in selection:
                current_margin = margin
    
                current_margin = current_margin * 0.5 if node.hide else current_margin  # use a smaller margin for hidden nodes
    
    
                if horizontal:
                    node.location.x = current_pos
                    current_pos += current_margin + node.dimensions.x
                    node.location.y = mid_y + (node.dimensions.y / 2)
                else:
                    node.location.y = current_pos
    
                    current_pos -= (current_margin * 0.3) + node.dimensions.y  # use half-margin for vertical alignment
    
                    node.location.x = mid_x - (node.dimensions.x / 2)
    
    
            # If active node is selected, center nodes around it
            if active_loc is not None:
                active_loc_diff = active_loc - nodes.active.location
                for node in selection:
                    node.location += active_loc_diff
            else:  # Position nodes centered around where they used to be
                locs = ([n.location.x + (n.dimensions.x / 2) for n in selection]) if horizontal else ([n.location.y - (n.dimensions.y / 2) for n in selection])
                new_mid = (max(locs) + min(locs)) / 2
                for node in selection:
                    if horizontal:
                        node.location.x += (mid_x - new_mid)
                    else:
                        node.location.y += (mid_y - new_mid)
    
    class NWSelectParentChildren(Operator, NWBase):
        bl_idname = "node.nw_select_parent_child"
    
        bl_label = "Select Parent or Children"
        bl_options = {'REGISTER', 'UNDO'}
    
    
    Campbell Barton's avatar
    Campbell Barton committed
        option: EnumProperty(
    
            name="option",
            items=(
                ('PARENT', 'Select Parent', 'Select Parent Frame'),
                ('CHILD', 'Select Children', 'Select members of selected frame'),
            )
        )
    
    
        def execute(self, context):
            nodes, links = get_nodes_links(context)
            option = self.option
            selected = [node for node in nodes if node.select]
            if option == 'PARENT':
                for sel in selected:
                    parent = sel.parent
                    if parent:
                        parent.select = True
            else:  # option == 'CHILD'
                for sel in selected:
                    children = [node for node in nodes if node.parent == sel]
                    for kid in children:
                        kid.select = True
    
            return {'FINISHED'}
    
    
    
    class NWDetachOutputs(Operator, NWBase):
    
        """Detach outputs of selected node leaving inputs linked"""
    
        bl_idname = "node.nw_detach_outputs"
    
        bl_label = "Detach Outputs"
    
    Bartek Skorupa's avatar
    Bartek Skorupa committed
        bl_options = {'REGISTER', 'UNDO'}
    
    Bartek Skorupa's avatar
    Bartek Skorupa committed
        def execute(self, context):
            nodes, links = get_nodes_links(context)
            selected = context.selected_nodes
            bpy.ops.node.duplicate_move_keep_inputs()
            new_nodes = context.selected_nodes
            bpy.ops.node.select_all(action="DESELECT")
            for node in selected:
                node.select = True
            bpy.ops.node.delete_reconnect()
            for new_node in new_nodes:
                new_node.select = True
    
            bpy.ops.transform.translate('INVOKE_DEFAULT')
    
    Bartek Skorupa's avatar
    Bartek Skorupa committed
            return {'FINISHED'}
    
    
    class NWLinkToOutputNode(Operator):
    
        """Link to Composite node or Material Output node"""
        bl_idname = "node.nw_link_out"
        bl_label = "Connect to Output"
    
    Bartek Skorupa's avatar
    Bartek Skorupa committed
        bl_options = {'REGISTER', 'UNDO'}
    
    Bartek Skorupa's avatar
    Bartek Skorupa committed
        @classmethod
        def poll(cls, context):
    
                if context.active_node is not None:
    
                    for out in context.active_node.outputs:
    
                        if is_visible_socket(out):
    
    Bartek Skorupa's avatar
    Bartek Skorupa committed
        def execute(self, context):
            nodes, links = get_nodes_links(context)
            active = nodes.active
    
            tree_type = context.space_data.tree_type
    
            shader_outputs = {'OBJECT':    'ShaderNodeOutputMaterial',
                              'WORLD':     'ShaderNodeOutputWorld',
                              'LINESTYLE': 'ShaderNodeOutputLineStyle'}
            output_type = {
                'ShaderNodeTree': shader_outputs[context.space_data.shader_type],
                'CompositorNodeTree': 'CompositorNodeComposite',
                'TextureNodeTree': 'TextureNodeOutput',
                'GeometryNodeTree': 'NodeGroupOutput',
            }[tree_type]
    
    Bartek Skorupa's avatar
    Bartek Skorupa committed
            for node in nodes:
    
                # check whether the node is an output node and,
                # if supported, whether it's the active one
                if node.rna_type.identifier == output_type \
                   and (node.is_active_output if hasattr(node, 'is_active_output')
                        else True):
    
    Bartek Skorupa's avatar
    Bartek Skorupa committed
                    output_node = node
                    break
    
            else:  # No output node exists
    
    Bartek Skorupa's avatar
    Bartek Skorupa committed
                bpy.ops.node.select_all(action="DESELECT")
    
                output_node = nodes.new(output_type)
    
                output_node.location.x = active.location.x + active.dimensions.x + 80
                output_node.location.y = active.location.y
    
    Bartek Skorupa's avatar
    Bartek Skorupa committed
                for i, output in enumerate(active.outputs):
    
                    if is_visible_socket(output):
    
                        output_index = i
                        break
                for i, output in enumerate(active.outputs):
    
                    if output.type == output_node.inputs[0].type and is_visible_socket(output):
    
    Bartek Skorupa's avatar
    Bartek Skorupa committed
                        output_index = i
                        break
    
    Greg Zaal's avatar
    Greg Zaal committed
                    if active.outputs[output_index].name == 'Volume':
                        out_input_index = 1
    
                    elif active.outputs[output_index].name == 'Displacement':
    
                elif tree_type == 'GeometryNodeTree':
                    if active.outputs[output_index].type != 'GEOMETRY':
                        return {'CANCELLED'}
    
                links.new(active.outputs[output_index], output_node.inputs[out_input_index])
    
    
            force_update(context)  # viewport render does not update
    
    Bartek Skorupa's avatar
    Bartek Skorupa committed
    
            return {'FINISHED'}
    
    
    
    class NWMakeLink(Operator, NWBase):
        """Make a link from one socket to another"""
        bl_idname = 'node.nw_make_link'
        bl_label = 'Make Link'
        bl_options = {'REGISTER', 'UNDO'}
    
    Campbell Barton's avatar
    Campbell Barton committed
        from_socket: IntProperty()
        to_socket: IntProperty()
    
    
        def execute(self, context):
            nodes, links = get_nodes_links(context)
    
            n1 = nodes[context.scene.NWLazySource]
            n2 = nodes[context.scene.NWLazyTarget]
    
            links.new(n1.outputs[self.from_socket], n2.inputs[self.to_socket])
    
    
            force_update(context)
    
    
            return {'FINISHED'}
    
    
    class NWCallInputsMenu(Operator, NWBase):
        """Link from this output"""
        bl_idname = 'node.nw_call_inputs_menu'
        bl_label = 'Make Link'
        bl_options = {'REGISTER', 'UNDO'}
    
    Campbell Barton's avatar
    Campbell Barton committed
        from_socket: IntProperty()
    
    
        def execute(self, context):
            nodes, links = get_nodes_links(context)
    
            context.scene.NWSourceSocket = self.from_socket
    
            n1 = nodes[context.scene.NWLazySource]
            n2 = nodes[context.scene.NWLazyTarget]
            if len(n2.inputs) > 1:
                bpy.ops.wm.call_menu("INVOKE_DEFAULT", name=NWConnectionListInputs.bl_idname)
            elif len(n2.inputs) == 1:
                links.new(n1.outputs[self.from_socket], n2.inputs[0])
            return {'FINISHED'}
    
    
    
    class NWAddSequence(Operator, NWBase, ImportHelper):
    
        """Add an Image Sequence"""
        bl_idname = 'node.nw_add_sequence'
        bl_label = 'Import Image Sequence'
        bl_options = {'REGISTER', 'UNDO'}
    
    
        directory: StringProperty(
            subtype="DIR_PATH"
        )
        filename: StringProperty(
            subtype="FILE_NAME"
        )
        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',
            default=True
        )
    
        def draw(self, context):
            layout = self.layout
            layout.alignment = 'LEFT'
    
            layout.prop(self, 'relative_path')
    
    
        def execute(self, context):
            nodes, links = get_nodes_links(context)
            directory = self.directory
            filename = self.filename
    
    Greg Zaal's avatar
    Greg Zaal committed
            files = self.files
            tree = context.space_data.node_tree
    
    Greg Zaal's avatar
    Greg Zaal committed
            # DEBUG
            # print ("\nDIR:", directory)
            # print ("FN:", filename)
            # print ("Fs:", list(f.name for f in files), '\n')
    
    Greg Zaal's avatar
    Greg Zaal committed
            if tree.type == 'SHADER':
    
                node_type = "ShaderNodeTexImage"
    
    Greg Zaal's avatar
    Greg Zaal committed
            elif tree.type == 'COMPOSITING':
    
                node_type = "CompositorNodeImage"
            else:
                self.report({'ERROR'}, "Unsupported Node Tree type!")
                return {'CANCELLED'}
    
    
    Greg Zaal's avatar
    Greg Zaal committed
            if not files[0].name and not filename:
                self.report({'ERROR'}, "No file chosen")
                return {'CANCELLED'}
            elif files[0].name and (not filename or not path.exists(directory+filename)):
                # User has selected multiple files without an active one, or the active one is non-existant
                filename = files[0].name
    
            if not path.exists(directory+filename):
                self.report({'ERROR'}, filename+" does not exist!")
                return {'CANCELLED'}
    
    
            without_ext = '.'.join(filename.split('.')[:-1])
    
    
            # if last digit isn't a number, it's not a sequence
    
    Greg Zaal's avatar
    Greg Zaal committed
            if not without_ext[-1].isdigit():
    
                self.report({'ERROR'}, filename+" does not seem to be part of a sequence")
                return {'CANCELLED'}
    
    
    
            extension = filename.split('.')[-1]
    
            reverse = without_ext[::-1] # reverse string
    
            count_numbers = 0
            for char in reverse:
    
                if char.isdigit():
    
                    count_numbers += 1
                else:
    
            without_num = without_ext[:count_numbers*-1]
    
            files = sorted(glob(directory + without_num + "[0-9]"*count_numbers + "." + extension))
    
            num_frames = len(files)
    
    
            nodes_list = [node for node in nodes]
            if nodes_list:
    
                nodes_list.sort(key=lambda k: k.location.x)
    
                xloc = nodes_list[0].location.x - 220  # place new nodes at far left
                yloc = 0
                for node in nodes:
                    node.select = False
                    yloc += node_mid_pt(node, 'y')
                yloc = yloc/len(nodes)
            else:
                xloc = 0
                yloc = 0
    
    
            name_with_hashes = without_num + "#"*count_numbers + '.' + extension
    
    
    Greg Zaal's avatar
    Greg Zaal committed
            bpy.ops.node.add_node('INVOKE_DEFAULT', use_transform=True, type=node_type)
    
            node = nodes.active
    
            node.label = name_with_hashes
    
            filepath = directory+(without_ext+'.'+extension)
            if self.relative_path:
                if bpy.data.filepath:
                    try:
                        filepath = bpy.path.relpath(filepath)
                    except ValueError:
                        pass
    
            img = bpy.data.images.load(filepath)
    
            img.source = 'SEQUENCE'
    
            img.name = name_with_hashes
    
            node.image = img
    
    Greg Zaal's avatar
    Greg Zaal committed
            image_user = node.image_user if tree.type == 'SHADER' else node
            image_user.frame_offset = int(files[0][len(without_num)+len(directory):-1*(len(extension)+1)]) - 1  # separate the number from the file name of the first  file
            image_user.frame_duration = num_frames
    
    class NWAddMultipleImages(Operator, NWBase, ImportHelper):
    
        """Add multiple images at once"""
        bl_idname = 'node.nw_add_multiple_images'
        bl_label = 'Open Selected Images'
        bl_options = {'REGISTER', 'UNDO'}
    
        directory: StringProperty(
            subtype="DIR_PATH"
        )
        files: CollectionProperty(
            type=bpy.types.OperatorFileListElement,
            options={'HIDDEN', 'SKIP_SAVE'}
        )
    
    
        def execute(self, context):
            nodes, links = get_nodes_links(context)
    
            xloc, yloc = context.region.view2d.region_to_view(context.area.width/2, context.area.height/2)
    
    
            if context.space_data.node_tree.type == 'SHADER':
                node_type = "ShaderNodeTexImage"
            elif context.space_data.node_tree.type == 'COMPOSITING':
                node_type = "CompositorNodeImage"
            else:
                self.report({'ERROR'}, "Unsupported Node Tree type!")
                return {'CANCELLED'}
    
            new_nodes = []
            for f in self.files:
                fname = f.name
    
                node = nodes.new(node_type)
                new_nodes.append(node)
                node.label = fname
                node.hide = True
                node.width_hidden = 100
                node.location.x = xloc
                node.location.y = yloc
                yloc -= 40
    
                img = bpy.data.images.load(self.directory+fname)
                node.image = img
    
            # shift new nodes up to center of tree
            list_size = new_nodes[0].location.y - new_nodes[-1].location.y
    
            for node in nodes:
                if node in new_nodes:
                    node.select = True
                    node.location.y += (list_size/2)
                else:
                    node.select = False
    
    class NWViewerFocus(bpy.types.Operator):
        """Set the viewer tile center to the mouse position"""
        bl_idname = "node.nw_viewer_focus"
        bl_label = "Viewer Focus"
    
    
    Campbell Barton's avatar
    Campbell Barton committed
        x: bpy.props.IntProperty()
        y: bpy.props.IntProperty()
    
    
        @classmethod
        def poll(cls, context):
            return nw_check(context) and context.space_data.tree_type == 'CompositorNodeTree'
    
        def execute(self, context):
            return {'FINISHED'}
    
        def invoke(self, context, event):
            render = context.scene.render
            space = context.space_data
            percent = render.resolution_percentage*0.01
    
    
            nodes, links = get_nodes_links(context)
            viewers = [n for n in nodes if n.type == 'VIEWER']
    
            if viewers:
                mlocx = event.mouse_region_x
                mlocy = event.mouse_region_y
    
                select_node = bpy.ops.node.select(location=(mlocx, mlocy), extend=False)
    
                if not 'FINISHED' in select_node:  # only run if we're not clicking on a node
                    region_x = context.region.width
                    region_y = context.region.height
    
                    region_center_x = context.region.width  / 2
                    region_center_y = context.region.height / 2
    
                    bd_x = render.resolution_x * percent * space.backdrop_zoom
                    bd_y = render.resolution_y * percent * space.backdrop_zoom
    
                    backdrop_center_x = (bd_x / 2) - space.backdrop_offset[0]
                    backdrop_center_y = (bd_y / 2) - space.backdrop_offset[1]
    
                    margin_x = region_center_x - backdrop_center_x
                    margin_y = region_center_y - backdrop_center_y
    
                    abs_mouse_x = (mlocx - margin_x) / bd_x
                    abs_mouse_y = (mlocy - margin_y) / bd_y
    
                    for node in viewers:
                        node.center_x = abs_mouse_x
                        node.center_y = abs_mouse_y
                else:
                    return {'PASS_THROUGH'}
    
    class NWSaveViewer(bpy.types.Operator, ExportHelper):
        """Save the current viewer node to an image file"""
        bl_idname = "node.nw_save_viewer"
        bl_label = "Save This Image"
    
    Campbell Barton's avatar
    Campbell Barton committed
        filepath: StringProperty(subtype="FILE_PATH")
        filename_ext: EnumProperty(
    
                name="Format",
                description="Choose the file format to save to",
    
                items=(('.bmp', "BMP", ""),
    
                       ('.rgb', 'IRIS', ""),
                       ('.png', 'PNG', ""),
                       ('.jpg', 'JPEG', ""),
                       ('.jp2', 'JPEG2000', ""),
                       ('.tga', 'TARGA', ""),
                       ('.cin', 'CINEON', ""),
                       ('.dpx', 'DPX', ""),
                       ('.exr', 'OPEN_EXR', ""),
                       ('.hdr', 'HDR', ""),
                       ('.tif', 'TIFF', "")),
                default='.png',
                )
    
        @classmethod
        def poll(cls, context):
    
            valid = False
            if nw_check(context):
                if context.space_data.tree_type == 'CompositorNodeTree':
                    if "Viewer Node" in [i.name for i in bpy.data.images]:
                        if sum(bpy.data.images["Viewer Node"].size) > 0:  # False if not connected or connected but no image
                            valid = True
            return valid
    
    
        def execute(self, context):
            fp = self.filepath
            if fp:
                formats = {
                           '.bmp': 'BMP',
                           '.rgb': 'IRIS',
                           '.png': 'PNG',
                           '.jpg': 'JPEG',
                           '.jpeg': 'JPEG',
                           '.jp2': 'JPEG2000',
                           '.tga': 'TARGA',
                           '.cin': 'CINEON',
                           '.dpx': 'DPX',
                           '.exr': 'OPEN_EXR',
                           '.hdr': 'HDR',
                           '.tiff': 'TIFF',
                           '.tif': 'TIFF'}
                basename, ext = path.splitext(fp)
                old_render_format = context.scene.render.image_settings.file_format
                context.scene.render.image_settings.file_format = formats[self.filename_ext]
                context.area.type = "IMAGE_EDITOR"
                context.area.spaces[0].image = bpy.data.images['Viewer Node']
                context.area.spaces[0].image.save_render(fp)
                context.area.type = "NODE_EDITOR"
                context.scene.render.image_settings.file_format = old_render_format
                return {'FINISHED'}
    
    
    
    class NWResetNodes(bpy.types.Operator):
        """Reset Nodes in Selection"""
        bl_idname = "node.nw_reset_nodes"
        bl_label = "Reset Nodes"
        bl_options = {'REGISTER', 'UNDO'}
    
        @classmethod
        def poll(cls, context):
            space = context.space_data
            return space.type == 'NODE_EDITOR'
    
        def execute(self, context):
            node_active = context.active_node
            node_selected = context.selected_nodes
            node_ignore = ["FRAME","REROUTE", "GROUP"]
    
            # Check if one node is selected at least
            if not (len(node_selected) > 0):
                self.report({'ERROR'}, "1 node must be selected at least")
                return {'CANCELLED'}
    
            active_node_name = node_active.name if node_active.select else None
            valid_nodes = [n for n in node_selected if n.type not in node_ignore]
    
            # Create output lists
            selected_node_names = [n.name for n in node_selected]
            success_names = []
    
            # Reset all valid children in a frame
            node_active_is_frame = False
            if len(node_selected) == 1 and node_active.type == "FRAME":
                node_tree = node_active.id_data
                children = [n for n in node_tree.nodes if n.parent == node_active]
                if children:
                    valid_nodes = [n for n in children if n.type not in node_ignore]
                    selected_node_names = [n.name for n in children if n.type not in node_ignore]
                    node_active_is_frame = True
    
            # Check if valid nodes in selection
            if not (len(valid_nodes) > 0):
                # Check for frames only
                frames_selected = [n for n in node_selected if n.type == "FRAME"]
                if (len(frames_selected) > 1 and len(frames_selected) == len(node_selected)):
                    self.report({'ERROR'}, "Please select only 1 frame to reset")
                else:
                    self.report({'ERROR'}, "No valid node(s) in selection")
                return {'CANCELLED'}
    
            # Report nodes that are not valid
            if len(valid_nodes) != len(node_selected) and node_active_is_frame is False:
                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 {}".format(", ".join(not_valid_names)))
    
            # Deselect all nodes
            for i in node_selected:
                i.select = False
    
            # Run through all valid nodes
    
            for node in valid_nodes:
    
    
                parent = node.parent if node.parent else None
                node_loc = [node.location.x, node.location.y]
    
                node_tree = node.id_data
                props_to_copy = 'bl_idname name location height width'.split(' ')
    
                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()])
    
                props = {j: getattr(node, j) for j in props_to_copy}
    
                new_node = node_tree.nodes.new(props['bl_idname'])
                props_to_copy.pop(0)
    
                for prop in props_to_copy:
                    setattr(new_node, prop, props[prop])
    
                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))
    
                new_node.select = False
                success_names.append(new_node.name)
    
            # Reselect all nodes
            if selected_node_names and node_active_is_frame is False:
                for i in selected_node_names:
                    node_tree.nodes[i].select = True
    
            if active_node_name is not None:
                node_tree.nodes[active_node_name].select = True
                node_tree.nodes.active = node_tree.nodes[active_node_name]
    
            self.report({'INFO'}, "Successfully reset {}".format(", ".join(success_names)))
            return {'FINISHED'}
    
    
    
    def drawlayout(context, layout, mode='non-panel'):
        tree_type = context.space_data.tree_type
    
        col = layout.column(align=True)
        col.menu(NWMergeNodesMenu.bl_idname)
        col.separator()
    
        col = layout.column(align=True)
        col.menu(NWSwitchNodeTypeMenu.bl_idname, text="Switch Node Type")
        col.separator()
    
    
            col = layout.column(align=True)
            col.operator(NWAddTextureSetup.bl_idname, text="Add Texture Setup", icon='NODE_SEL')
    
            col.operator(NWAddPrincipledSetup.bl_idname, text="Add Principled Setup", icon='NODE_SEL')
    
            col.separator()
    
        col = layout.column(align=True)
        col.operator(NWDetachOutputs.bl_idname, icon='UNLINKED')
    
        col.operator(NWSwapLinks.bl_idname)
    
        col.menu(NWAddReroutesMenu.bl_idname, text="Add Reroutes", icon='LAYER_USED')
        col.separator()
    
        col = layout.column(align=True)
        col.menu(NWLinkActiveToSelectedMenu.bl_idname, text="Link Active To Selected", icon='LINKED')
    
        if tree_type != 'GeometryNodeTree':
            col.operator(NWLinkToOutputNode.bl_idname, icon='DRIVER')
    
        col.separator()
    
        col = layout.column(align=True)
        if mode == 'panel':
            row = col.row(align=True)
            row.operator(NWClearLabel.bl_idname).option = True
            row.operator(NWModifyLabels.bl_idname)
        else:
            col.operator(NWClearLabel.bl_idname).option = True
            col.operator(NWModifyLabels.bl_idname)
        col.menu(NWBatchChangeNodesMenu.bl_idname, text="Batch Change")
        col.separator()
        col.menu(NWCopyToSelectedMenu.bl_idname, text="Copy to Selected")
        col.separator()
    
        col = layout.column(align=True)
        if tree_type == 'CompositorNodeTree':
            col.operator(NWResetBG.bl_idname, icon='ZOOM_PREVIOUS')
    
        if tree_type != 'GeometryNodeTree':
            col.operator(NWReloadImages.bl_idname, icon='FILE_REFRESH')
    
        col.separator()
    
        col = layout.column(align=True)
        col.operator(NWFrameSelected.bl_idname, icon='STICKY_UVS_LOC')
        col.separator()
    
    
        col = layout.column(align=True)
    
        col.operator(NWAlignNodes.bl_idname, icon='CENTER_ONLY')
    
        col = layout.column(align=True)
        col.operator(NWDeleteUnused.bl_idname, icon='CANCEL')
        col.separator()
    
    
    class NodeWranglerPanel(Panel, NWBase):
        bl_idname = "NODE_PT_nw_node_wrangler"
    
        bl_space_type = 'NODE_EDITOR'
    
        bl_category = "Node Wrangler"
    
    Campbell Barton's avatar
    Campbell Barton committed
        prepend: StringProperty(
    
    Campbell Barton's avatar
    Campbell Barton committed
        append: StringProperty()
        remove: StringProperty()
    
    
        def draw(self, context):
    
            self.layout.label(text="(Quick access: Shift+W)")
    
            drawlayout(context, self.layout, mode='panel')
    
    #
    #  M E N U S
    #
    class NodeWranglerMenu(Menu, NWBase):
        bl_idname = "NODE_MT_nw_node_wrangler_menu"
        bl_label = "Node Wrangler"
    
    
        def draw(self, context):
    
            self.layout.operator_context = 'INVOKE_DEFAULT'
    
            drawlayout(context, self.layout)
    
    
    class NWMergeNodesMenu(Menu, NWBase):
        bl_idname = "NODE_MT_nw_merge_nodes_menu"
    
        bl_label = "Merge Selected Nodes"
    
        def draw(self, context):
            type = context.space_data.tree_type
            layout = self.layout
    
                layout.menu(NWMergeShadersMenu.bl_idname, text="Use Shaders")
    
            if type == 'GeometryNodeTree':
                layout.menu(NWMergeGeometryMenu.bl_idname, text="Use Geometry Nodes")
                layout.menu(NWMergeMathMenu.bl_idname, text="Use Math Nodes")
            else:
                layout.menu(NWMergeMixMenu.bl_idname, text="Use Mix Nodes")
                layout.menu(NWMergeMathMenu.bl_idname, text="Use Math Nodes")
                props = layout.operator(NWMergeNodes.bl_idname, text="Use Z-Combine Nodes")
                props.mode = 'MIX'
                props.merge_type = 'ZCOMBINE'
                props = layout.operator(NWMergeNodes.bl_idname, text="Use Alpha Over Nodes")
                props.mode = 'MIX'
                props.merge_type = 'ALPHAOVER'
    
    class NWMergeGeometryMenu(Menu, NWBase):
        bl_idname = "NODE_MT_nw_merge_geometry_menu"
        bl_label = "Merge Selected Nodes using Geometry Nodes"
        def draw(self, context):
            layout = self.layout
            # The boolean node + Join Geometry node
            for type, name, description in geo_combine_operations:
                props = layout.operator(NWMergeNodes.bl_idname, text=name)
                props.mode = type
                props.merge_type = 'GEOMETRY'
    
    class NWMergeShadersMenu(Menu, NWBase):
        bl_idname = "NODE_MT_nw_merge_shaders_menu"
    
        bl_label = "Merge Selected Nodes using Shaders"
    
        def draw(self, context):
            layout = self.layout
    
            for type in ('MIX', 'ADD'):
                props = layout.operator(NWMergeNodes.bl_idname, text=type)
    
                props.mode = type
                props.merge_type = 'SHADER'
    
    
    
    class NWMergeMixMenu(Menu, NWBase):
        bl_idname = "NODE_MT_nw_merge_mix_menu"
    
        bl_label = "Merge Selected Nodes using Mix"
    
        def draw(self, context):
            layout = self.layout
            for type, name, description in blend_types:
    
                props = layout.operator(NWMergeNodes.bl_idname, text=name)
    
                props.mode = type
                props.merge_type = 'MIX'
    
    
    
    class NWConnectionListOutputs(Menu, NWBase):
        bl_idname = "NODE_MT_nw_connection_list_out"
        bl_label = "From:"
    
        def draw(self, context):
            layout = self.layout
            nodes, links = get_nodes_links(context)
    
            n1 = nodes[context.scene.NWLazySource]
    
            for index, output in enumerate(n1.outputs):
    
                # Only show sockets that are exposed.
    
                if output.enabled:
                    layout.operator(NWCallInputsMenu.bl_idname, text=output.name, icon="RADIOBUT_OFF").from_socket=index
    
    
    
    class NWConnectionListInputs(Menu, NWBase):
        bl_idname = "NODE_MT_nw_connection_list_in"
        bl_label = "To:"
    
        def draw(self, context):
            layout = self.layout
            nodes, links = get_nodes_links(context)
    
            n2 = nodes[context.scene.NWLazyTarget]
    
    
            for index, input in enumerate(n2.inputs):
    
                # Only show sockets that are exposed.
                # This prevents, for example, the scale value socket
                # of the vector math node being added to the list when
    
                # the mode is not 'SCALE'.
    
                if input.enabled:
                    op = layout.operator(NWMakeLink.bl_idname, text=input.name, icon="FORWARD")
    
                    op.from_socket = context.scene.NWSourceSocket
                    op.to_socket = index
    
    class NWMergeMathMenu(Menu, NWBase):
        bl_idname = "NODE_MT_nw_merge_math_menu"
    
        bl_label = "Merge Selected Nodes using Math"
    
        def draw(self, context):
            layout = self.layout
            for type, name, description in operations:
    
                props = layout.operator(NWMergeNodes.bl_idname, text=name)
    
                props.mode = type
                props.merge_type = 'MATH'
    
    
    
    class NWBatchChangeNodesMenu(Menu, NWBase):
        bl_idname = "NODE_MT_nw_batch_change_nodes_menu"
    
        bl_label = "Batch Change Selected Nodes"
    
        def draw(self, context):
            layout = self.layout
    
            layout.menu(NWBatchChangeBlendTypeMenu.bl_idname)
            layout.menu(NWBatchChangeOperationMenu.bl_idname)
    
    class NWBatchChangeBlendTypeMenu(Menu, NWBase):
        bl_idname = "NODE_MT_nw_batch_change_blend_type_menu"
    
        bl_label = "Batch Change Blend Type"
    
        def draw(self, context):
            layout = self.layout
            for type, name, description in blend_types:
    
                props = layout.operator(NWBatchChangeNodes.bl_idname, text=name)
    
                props.blend_type = type
                props.operation = 'CURRENT'
    
    
    
    class NWBatchChangeOperationMenu(Menu, NWBase):
        bl_idname = "NODE_MT_nw_batch_change_operation_menu"
    
        bl_label = "Batch Change Math Operation"
    
        def draw(self, context):
            layout = self.layout
            for type, name, description in operations:
    
                props = layout.operator(NWBatchChangeNodes.bl_idname, text=name)
    
                props.blend_type = 'CURRENT'
                props.operation = type