ui.py

# SPDX-License-Identifier: GPL-2.0-or-later

import bpy
from bpy.props import (
    BoolProperty,
    IntProperty,
    EnumProperty,
    StringProperty
)

from mathutils import Color

from .utils.errors import MetarigError
from .utils.rig import write_metarig
from .utils.widgets import write_widget
from .utils.naming import unique_name
from .utils.rig import upgradeMetarigTypes, outdated_types

from .rigs.utils import get_limb_generated_names

from .utils.animation import get_keyed_frames_in_range, bones_in_frame, overwrite_prop_animation
from .utils.animation import RIGIFY_OT_get_frame_range

from .utils.animation import register as animation_register
from .utils.animation import unregister as animation_unregister

from . import base_rig
from . import rig_lists
from . import generate
from . import rot_mode
from . import feature_set_list


def build_type_list(context, rigify_types):
    rigify_types.clear()

    for r in sorted(rig_lists.rigs):
        if (context.object.data.active_feature_set in ('all', rig_lists.rigs[r]['feature_set'])
                or len(feature_set_list.get_enabled_modules_names()) == 0
                ):
            a = rigify_types.add()
            a.name = r


class DATA_PT_rigify(bpy.types.Panel):
    bl_label = "Rigify"
    bl_space_type = 'PROPERTIES'
    bl_region_type = 'WINDOW'
    bl_context = "data"

    @classmethod
    def poll(cls, context):
        obj = context.object
        if not context.object:
            return False
        return obj.type == 'ARMATURE' \
            and obj.data.get("rig_id") is None

    def draw(self, context):
        C = context
        layout = self.layout
        obj = C.object

        WARNING = "Warning: Some features may change after generation"
        show_warning = False
        show_update_metarig = False
        show_not_updatable = False
        show_upgrade_face = False

        check_props = ['IK_follow', 'root/parent', 'FK_limb_follow', 'IK_Stretch']

        for posebone in obj.pose.bones:
            bone = posebone.bone
            if not bone:
                # If we are in edit mode and the bone was just created,
                # a pose bone won't exist yet.
                continue
            if bone.layers[30] and (list(set(posebone.keys()) & set(check_props))):
                show_warning = True
                break

        for b in obj.pose.bones:
            if b.rigify_type in outdated_types.keys():
                old_bone = b.name
                old_rig = b.rigify_type
                if outdated_types[b.rigify_type]:
                    show_update_metarig = True
                else:
                    show_update_metarig = False
                    show_not_updatable = True
                    break
            elif b.rigify_type == 'faces.super_face':
                show_upgrade_face = True

        if show_warning:
            layout.label(text=WARNING, icon='ERROR')

        enable_generate = not (show_not_updatable or show_update_metarig)

        if show_not_updatable:
            layout.label(text="WARNING: This metarig contains deprecated rigify rig-types and cannot be upgraded automatically.", icon='ERROR')
            layout.label(text="("+old_rig+" on bone "+old_bone+")")
        elif show_update_metarig:
            layout.label(text="This metarig contains old rig-types that can be automatically upgraded to benefit of rigify's new features.", icon='ERROR')
            layout.label(text="("+old_rig+" on bone "+old_bone+")")
            layout.operator("pose.rigify_upgrade_types", text="Upgrade Metarig")
        elif show_upgrade_face:
            layout.label(text="This metarig uses the old face rig.", icon='INFO')
            layout.operator("pose.rigify_upgrade_face")

        row = layout.row()
        # Rig type field

        col = layout.column(align=True)
        col.active = (not 'rig_id' in C.object.data)

        col.separator()
        row = col.row()
        text = "Re-Generate Rig" if obj.data.rigify_target_rig else "Generate Rig"
        row.operator("pose.rigify_generate", text=text, icon='POSE_HLT')
        row.enabled = enable_generate


class DATA_PT_rigify_advanced(bpy.types.Panel):
    bl_space_type = 'PROPERTIES'
    bl_region_type = 'WINDOW'
    bl_context = "data"
    bl_label = "Advanced"
    bl_parent_id = 'DATA_PT_rigify'
    bl_options = {'DEFAULT_CLOSED'}

    def draw(self, context):
        layout = self.layout
        layout.use_property_split = True
        layout.use_property_decorate = False

        armature_id_store = context.object.data

        col = layout.column()
        col.row().prop(armature_id_store, "rigify_target_rig", text="Target Rig")
        col.row().prop(armature_id_store, "rigify_rig_ui", text="Rig UI Script")
        col.separator()
        col.row().prop(armature_id_store, "rigify_widgets_collection")
        col.row().prop(armature_id_store, "rigify_force_widget_update")
        col.row().prop(armature_id_store, "rigify_mirror_widgets")
        col.separator()
        col.row().prop(armature_id_store, "rigify_finalize_script", text="Run Script")


class DATA_PT_rigify_samples(bpy.types.Panel):
    bl_label = "Samples"
    bl_space_type = 'PROPERTIES'
    bl_region_type = 'WINDOW'
    bl_context = "data"
    bl_parent_id = "DATA_PT_rigify"
    bl_options = {'DEFAULT_CLOSED'}

    @classmethod
    def poll(cls, context):
        obj = context.object
        if not obj:
            return False
        return obj.type == 'ARMATURE' \
            and obj.data.get("rig_id") is None \
            and obj.mode == 'EDIT'

    def draw(self, context):
        layout = self.layout
        layout.use_property_split = True
        layout.use_property_decorate = False
        obj = context.object
        id_store = context.window_manager

        # Build types list
        build_type_list(context, id_store.rigify_types)

        if id_store.rigify_active_type > len(id_store.rigify_types):
            id_store.rigify_active_type = 0

        # Rig type list
        if len(feature_set_list.get_enabled_modules_names()) > 0:
            row = layout.row()
            row.prop(context.object.data, "active_feature_set")
        row = layout.row()
        row.template_list("UI_UL_list", "rigify_types", id_store, "rigify_types", id_store, 'rigify_active_type')

        props = layout.operator("armature.metarig_sample_add", text="Add sample")
        props.metarig_type = id_store.rigify_types[id_store.rigify_active_type].name


class DATA_PT_rigify_layer_names(bpy.types.Panel):
    bl_label = "Layer Names"
    bl_space_type = 'PROPERTIES'
    bl_region_type = 'WINDOW'
    bl_context = "data"
    bl_options = {'DEFAULT_CLOSED'}
    bl_parent_id = "DATA_PT_rigify"

    @classmethod
    def poll(cls, context):