rig_picker/area.py

import bpy
from bpy.types import NodeTree, NODE_PT_tools_active, NODE_HT_header, Menu
from .constants import PICKERS


# Derived from the NodeTree base type, similar to Menu, Operator, Panel, etc.
class RigPickerTree(NodeTree):
    # Description string
    '''A custom node tree type that will show up in the editor type list'''
    # Optional identifier string. If not explicitly defined, the python class name is used.
    # Label for nice name display
    bl_label = "Rig Picker"
    # Icon identifier
    bl_icon = 'OUTLINER_DATA_ARMATURE'


class RP_MT_picker(Menu):
    """Picker"""
    bl_label = "Picker"

    def draw(self, context):
        layout = self.layout
        scn = context.scene


        layout.operator('rigpicker.reload_picker', icon='FILE_REFRESH')
        row = layout.row(align=True)

        # Has at least one picker collection in the scene
        if not [c.rig_picker.enabled for c in scn.collection.children_recursive]:
            row.enabled = False
        row.prop(scn.rig_picker, 'use_pick_bone', text='Auto Bone Assign')


class RP_MT_animation(Menu):
    """Picker"""
    bl_label = "Animation"

    def draw(self, context):
        layout = self.layout

        layout.operator("rigpicker.toogle_bone_layer", icon="MOUSE_LMB_DRAG")

        row = layout.row()
        op = row.operator("node.context_menu_picker", icon="MOUSE_RMB")
        if not context.active_pose_bone:
            row.enabled = False

        layout.separator()
        op = layout.operator("rigpicker.call_operator", text='Select All')
        op.operator = "pose.select_all"
        op.arguments = '{"action": "SELECT"}'

        op = layout.operator("rigpicker.call_operator", text='Select All')
        op.operator = "pose.select_all"
        op.arguments = '{"action": "DESELECT"}'

        op = layout.operator("rigpicker.call_operator", text='Frame Selected')
        op.operator = "view3d.view_selected"
        op.view_3d = True

        layout.separator()
        layout.operator("rigpicker.call_operator", text='Insert Keyframe').operator="animtoolbox.insert_keyframe"
        layout.operator("anim.keyframe_delete_v3d", text='Delete Keyframe')

        layout.separator()
        op = layout.operator("rigpicker.call_operator", text='Move')
        op.operator="transform.translate"
        op.invoke = True
        op.view_3d = True
        
        layout.operator("node.picker_transform", text='Rotate').mode='ROTATE'
        layout.operator("node.picker_transform", text='Scale').mode='SCALE'
        layout.separator()
        layout.operator("rigpicker.call_operator", text='Reset Bone').operator="animtoolbox.reset_bone"
        layout.operator("rigpicker.call_operator", text='Clear Location').operator='pose.loc_clear'
        layout.operator("rigpicker.call_operator", text='Clear Rotation').operator='pose.rot_clear'
        layout.operator("rigpicker.call_operator", text='Clear Scale').operator='pose.scale_clear'

def draw_header(self, context):
    if not context.space_data.tree_type == 'RigPickerTree':
        self._draw(context)
        return
    
    scn = context.scene

    layout = self.layout
    layout.template_header()

    if not context.space_data.node_tree:
        ntree = bpy.data.node_groups.get('.rig_picker')
        if not ntree:
            ntree = bpy.data.node_groups.new('.rig_picker', 'RigPickerTree')
        
        context.space_data.node_tree = ntree

    row = layout.row(align=True)
    row.menu("RP_MT_picker")
    row.menu("RP_MT_animation")

    ob = context.object
    if not ob or not ob.type == 'ARMATURE':
        return
    
    picker_group = PICKERS.get(ob)
    if not picker_group:
        return

    if scn.rig_picker.use_pick_bone:
        layout.alert = True
        layout.label(text='Auto Bone Assign')
        layout.prop(scn.rig_picker, 'use_pick_bone', icon='PANEL_CLOSE', text='', emboss=False)

    layout.separator_spacer()
    row = layout.row()
    row.enabled = False
    row.label(text=ob.name)

    layout.separator_spacer()
    row = layout.row(align=True)

    for i, picker in enumerate(picker_group.pickers):
        row.operator('rigpicker.fit_picker', text=f'{i+1}').index=i

    row.operator('rigpicker.fit_picker', text='', icon='FULLSCREEN_ENTER').index = -1
        

def tools_from_context(context, mode=None):
    sp = context.space_data
    if sp and sp.type == 'NODE_EDITOR' and sp.tree_type == 'RigPickerTree':
        return []
    else:
        return NODE_PT_tools_active._tools_from_context(context, mode)

def tool_set_by_id(self, context):
    sd = context.space_data
    if sd and sd.type == 'NODE_EDITOR' and sd.tree_type == 'RigPickerTree':
        return {"FINISHED"}
    else:
        return self._execute(context)

@classmethod
def poll(cls, context):
    sp = context.space_data
    if sp and sp.type == 'NODE_EDITOR' and sp.tree_type == 'RigPickerTree':
        return False
    else:
        return self._poll(context)

classes = (
    RigPickerTree,
    RP_MT_picker,
    RP_MT_animation
)


def register():
    # Remove the tools inside the picker space
    bpy.types.WM_OT_tool_set_by_id._execute = bpy.types.WM_OT_tool_set_by_id.execute #tool_set_by_id    
    bpy.types.WM_OT_tool_set_by_id.execute = tool_set_by_id

    NODE_PT_tools_active._tools_from_context = NODE_PT_tools_active.tools_from_context
    NODE_PT_tools_active.tools_from_context = tools_from_context

    NODE_HT_header._draw = NODE_HT_header.draw
    NODE_HT_header.draw = draw_header

    for cls in classes:
        bpy.utils.register_class(cls)

def unregister():
    for cls in reversed(classes):
        bpy.utils.unregister_class(cls)
    
    bpy.types.WM_OT_tool_set_by_id.execute = bpy.types.WM_OT_tool_set_by_id._execute 
    del bpy.types.WM_OT_tool_set_by_id._execute

    NODE_PT_tools_active.tools_from_context = NODE_PT_tools_active._tools_from_context
    del NODE_PT_tools_active._tools_from_context

    NODE_HT_header.draw = NODE_HT_header._draw
    del NODE_HT_header._draw