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