Skip to content
Snippets Groups Projects
node_wrangler.py 183 KiB
Newer Older
  • Learn to ignore specific revisions
  •                             ('RGBA', [t[0] for t in blend_types], selected_mix),
                                ('VALUE', [t[0] for t in operations], selected_math),
    
                            output_type = node.outputs[0].type
                            valid_mode = mode in types_list
                            # 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.
                            if 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),
    
                                ('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_math, selected_z, selected_alphaover]:
    
                if nodes_list:
                    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)
    
                    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
                    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
    
                            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
    
                            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:
                                    loc_y = loc_y - 50
    
                                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:
                                    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
                    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]
    
                    # 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 marged are linked to one another.
                                # Create list of invalid indexes.
                                invalid_i = [n[0] for n in (selected_mix + selected_math + selected_shader + selected_z)]
                                # Link only if "to_node" index not in invalid indexes list.
                                if ss_link.to_node not in [nodes[i] for i in invalid_i]:
                                    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:
    
                        # Prevent cyclic dependencies when nodes to be marged are linked to one another.
                        # Create list of invalid indexes.
    
                        invalid_i = [n[0] for n in (selected_mix + selected_math + selected_shader + selected_z)]
    
                        # Link only if "to_node" index not in invalid indexes list.
                        if fs_link.to_node not in [nodes[i] for i in invalid_i]:
                            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
    
                    for i, x, y, dx, h in nodes_list:
    
                        nodes[i].select = False
    
            return {'FINISHED'}
    
    
    
    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'}
    
        blend_type = EnumProperty(
    
            name="Blend Type",
            items=blend_types + navs,
        )
    
        operation = EnumProperty(
    
            name="Operation",
            items=operations + navs,
        )
    
    
        def execute(self, context):
    
            nodes, links = get_nodes_links(context)
            blend_type = self.blend_type
            operation = self.operation
            for node in context.selected_nodes:
                if node.type == 'MIX_RGB':
                    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':
                    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.
        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 is not 'FRAME':
                    valid = True
    
            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'}
    
        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'}
    
        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'}
    
        prepend = StringProperty(
            name="Add to Beginning"
        )
        append = StringProperty(
            name="Add to End"
        )
        replace_from = StringProperty(
            name="Text to Replace"
        )
        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'}
    
    
        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):
            valid = False
    
            if nw_check(context):
                space = context.space_data
    
                if space.tree_type == 'ShaderNodeTree' and context.scene.render.engine == 'CYCLES':
    
                    valid = True
            return valid
    
        def execute(self, context):
            nodes, links = get_nodes_links(context)
    
            shader_types = [x[1] for x in shaders_shader_nodes_props if x[1] not in {'MIX_SHADER', 'ADD_SHADER'}]
            texture_types = [x[1] for x in shaders_texture_nodes_props]
    
            selected_nodes = [n for n in nodes if n.select]
            for t_node in selected_nodes:
                valid = False
                input_index = 0
                if t_node.inputs:
                    for index, i in enumerate(t_node.inputs):
                        if not i.is_linked:
    
                            input_index = index
                            break
                if valid:
                    locx = t_node.location.x
                    locy = t_node.location.y - t_node.dimensions.y/2
    
                    xoffset = [500, 700]
                    is_texture = False
                    if t_node.type in texture_types + ['MAPPING']:
                        xoffset = [290, 500]
                        is_texture = True
    
                    coordout = 2
                    image_type = 'ShaderNodeTexImage'
    
                    if (t_node.type in texture_types and t_node.type != 'TEX_IMAGE') or (t_node.type == 'BACKGROUND'):
                        coordout = 0  # image texture uses UVs, procedural textures and Background shader use Generated
                        if t_node.type == 'BACKGROUND':
                            image_type = 'ShaderNodeTexEnvironment'
    
                    if not is_texture:
                        tex = nodes.new(image_type)
                        tex.location = [locx - 200, locy + 112]
                        nodes.active = tex
                        links.new(tex.outputs[0], t_node.inputs[input_index])
    
                    t_node.select = False
                    if self.add_mapping or is_texture:
                        if t_node.type != 'MAPPING':
                            m = nodes.new('ShaderNodeMapping')
                            m.location = [locx - xoffset[0], locy + 141]
                            m.width = 240
                        else:
                            m = t_node
                        coord = nodes.new('ShaderNodeTexCoord')
                        coord.location = [locx - (200 if t_node.type == 'MAPPING' else xoffset[1]), locy + 124]
    
                        if not is_texture:
                            links.new(m.outputs[0], tex.inputs[0])
                            links.new(coord.outputs[coordout], m.inputs[0])
                        else:
                            nodes.active = m
                            links.new(m.outputs[0], t_node.inputs[input_index])
                            links.new(coord.outputs[coordout], m.inputs[0])
    
                    self.report({'WARNING'}, "No free inputs for node: "+t_node.name)
    
    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'}
    
        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()
                    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 is not '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
                            for render_pass, out_name, exr_name, in_internal, in_cycles in rl_outputs:
                                if output.name == out_name:
                                    pass_used = getattr(node_scene.render.layers[node_layer], 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'}
    
        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 render_pass, out_name, exr_name, in_internal, in_cycles 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 == out_name:
                            # example 'render_pass' entry: 'use_pass_uv' Check if True in scene render layers
                            pass_used = getattr(active.scene.render.layers[active.layer], 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 asignment 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 render_pass, out_name, exr_name, in_internal, in_cycles in rl_outputs:
                                if out.name in {out_name, exr_name}:
                                    src_name = (out_name, exr_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'}
    
        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 layed 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'}
    
        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 inluts liked"""
        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, NWBase):
        """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):
    
            valid = False
            if nw_check(context):
                if context.active_node is not None:
    
                    for out in context.active_node.outputs:
                        if not out.hide:
                            valid = True
                            break
    
    Bartek Skorupa's avatar
    Bartek Skorupa committed
        def execute(self, context):
            nodes, links = get_nodes_links(context)
            active = nodes.active
            output_node = None
    
            tree_type = context.space_data.tree_type
            output_types_shaders = [x[1] for x in shaders_output_nodes_props]
            output_types_compo = ['COMPOSITE']
    
            output_types_blender_mat = ['OUTPUT']
            output_types_textures = ['OUTPUT']
            output_types = output_types_shaders + output_types_compo + output_types_blender_mat
    
    Bartek Skorupa's avatar
    Bartek Skorupa committed
            for node in nodes:
    
                if node.type in output_types:
    
    Bartek Skorupa's avatar
    Bartek Skorupa committed
                    output_node = node
                    break
            if not output_node:
                bpy.ops.node.select_all(action="DESELECT")
    
                if tree_type == 'ShaderNodeTree':
    
                    if context.scene.render.engine == 'CYCLES':
                        output_node = nodes.new('ShaderNodeOutputMaterial')
                    else:
                        output_node = nodes.new('ShaderNodeOutput')
    
                elif tree_type == 'CompositorNodeTree':
    
    Bartek Skorupa's avatar
    Bartek Skorupa committed
                    output_node = nodes.new('CompositorNodeComposite')
    
                elif tree_type == 'TextureNodeTree':
                    output_node = nodes.new('TextureNodeOutput')
    
                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
            if (output_node and active.outputs):
                for i, output in enumerate(active.outputs):
    
                    if not output.hide:
                        output_index = i
                        break
                for i, output in enumerate(active.outputs):
                    if output.type == output_node.inputs[0].type and not output.hide:
    
    Bartek Skorupa's avatar
    Bartek Skorupa committed
                        output_index = i
                        break
    
                if tree_type == 'ShaderNodeTree' and context.scene.render.engine == 'CYCLES':
    
    Greg Zaal's avatar
    Greg Zaal committed
                    if active.outputs[output_index].name == 'Volume':
                        out_input_index = 1
                    elif active.outputs[output_index].type != 'SHADER':  # connect to displacement if not a shader
    
                        out_input_index = 2
                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'}
        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'}
        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, 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")
    
    Greg Zaal's avatar
    Greg Zaal committed
        files = CollectionProperty(type=bpy.types.OperatorFileListElement, options={'HIDDEN', 'SKIP_SAVE'})
    
    
        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')