asset_library/action/functions.py

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


"""
Functions related to anim and pose.
"""

from collections.abc import Collection

from typing import Optional, FrozenSet, Set, Union, Iterable, cast
import dataclasses
import functools
import re

import bpy
from bpy.types import (
    Action,
    Bone,
    Context,
    FCurve,
    Keyframe,
    Object,
    TimelineMarker
)

from asset_library.common.bl_utils import active_catalog_id, split_path

FCurveValue = Union[float, int]

pose_bone_re = re.compile(r'pose.bones\["([^"]+)"\]')
"""RegExp for matching FCurve data paths."""

def is_pose(action):
    for fc in action.fcurves:
        if len(fc.keyframe_points) > 1:
            return False

    return True

def get_bone_visibility(data_path):
    bone, prop = split_path(data_path)
    
    ob = bpy.context.object
    b_layers = [i for i, val in enumerate(ob.pose.bones[bone].bone.layers) if val]
    
    rig_layers = [(i, val) for i, val in enumerate(ob.data.layers)]

    return ob.data.layers[b_layers[0]]

def get_keyframes(action, selected=False, includes=[]):
    if selected:
        # keyframes = sorted([int(k.co[0]) for f in action.fcurves for k in f.keyframe_points if k.select_control_point and get_bone_visibility(f.data_path)])
        keyframes = []
        for f in action.fcurves:
            bone, prop = split_path(f.data_path)
            for k in f.keyframe_points:
                if bone not in includes:
                    continue
                if not k.select_control_point:
                    continue
                if not get_bone_visibility(f.data_path):
                    continue
                
                keyframes += [int(k.co[0])]                    
        if len(keyframes) <= 1:
            keyframes = [bpy.context.scene.frame_current]
    else:
        keyframes = sorted([int(k.co[0]) for f in action.fcurves for k in f.keyframe_points])

    return keyframes

def get_marker(action):
    if action.pose_markers:
        markers = action.pose_markers
        return next((m.name for m in markers if m.frame == bpy.context.scene.frame_current), None)

def reset_bone(bone, transform=True, custom_props=True):
    if transform:
        bone.location = (0, 0, 0)
        if bone.rotation_mode == "QUATERNION":
            bone.rotation_quaternion = (0, 0, 0, 0)
        elif bone.rotation_mode == "AXIS_ANGLE":
            bone.rotation_axis_angle = (0, 0, 0, 0)
        else:
            bone.rotation_euler = (0, 0, 0)
            
        bone.scale = (1, 1, 1)
    
    if custom_props:
        for key, value in bone.items():
            try:
                id_prop = bone.id_properties_ui(key)
            except TypeError:
                continue
                
            if not isinstance(value, (int, float)) or not id_prop:
                continue
            bone[key] = id_prop.as_dict()['default']

def is_asset_action(action):
    return action.asset_data and action.asset_data.catalog_id != str(uuid.UUID(int=0))

def conform_action(action):
    tags = ('pose', 'anim')

    if any(tag in action.asset_data.tags.keys() for tag in tags):
        return

    for fc in action.fcurves:
        action.asset_data['is_single_frame'] = True
        if len(fc.keyframe_points) > 1:
            action.asset_data['is_single_frame'] = False
            break

    if action.asset_data['is_single_frame']:
        action.asset_data.tags.new('pose')
    else:
        action.asset_data.tags.new('anim')

def clean_action(action='', frame_start=0, frame_end=0, excludes=[], includes=[]):
    ## Clean Keyframe Before/After Range
    for fc in action.fcurves:
        bone, prop = split_path(fc.data_path)
        # !! Mush Mush dependent. Need to be fix
        if bone in excludes or bone not in includes:
            action.fcurves.remove(fc)
            continue

        # Add Keyframe At Start/End Range
        for fr in (frame_start, frame_end):
            fc_val = fc.evaluate(fr)
            fc.keyframe_points.insert(frame=fr, value=fc_val)
            fc.update()

        # Remove Keyframe out of range
        for k in reversed(fc.keyframe_points):                    
            if int(k.co[0]) not in range(frame_start, frame_end+1):
                fc.keyframe_points.remove(k)
                fc.update()

def append_action(action_path='', action_name=''):
    print(f'Loading {action_name} from: {action_path}')

    with bpy.data.libraries.load(str(action_path), link=False) as (data_from, data_to):
        data_to.actions = [action_name]

    return data_to.actions[0]

def apply_anim(action_lib, ob, bones=[]):
    from mathutils import Vector

    scn = bpy.context.scene

    if not ob.animation_data:
        ob.animation_data_create()
    
    action = ob.animation_data.action
    
    if not action:
        action = bpy.data.actions.new(ob.name)
        ob.animation_data.action = action

    keys = sorted([k.co[0] for f in action_lib.fcurves for k in f.keyframe_points])
    if not keys:
        print(f'The action {action_lib.name} has no keyframes')
        return

    first_key = keys[0]
    key_offset = scn.frame_current - first_key

    key_attr = ('type', 'interpolation', 'handle_left_type', 'handle_right_type',
        'amplitude', 'back', 'easing', 'period', 'handle_right', 'handle_left'
    )
    for fc in action_lib.fcurves:
        bone_name, prop_name = split_path(fc.data_path)

        if bones and bone_name not in bones:
            continue

        action_fc = action.fcurves.find(fc.data_path, index=fc.array_index)
        if not action_fc:
            action_fc = action.fcurves.new(
                fc.data_path,
                index=fc.array_index,
                action_group=fc.group.name if fc.group else fc.data_path.split('"')[1]
            )
        
        for kf_lib in fc.keyframe_points:
            kf = action_fc.keyframe_points.insert(
                frame=kf_lib.co[0] + key_offset,
                value=kf_lib.co[1]
            )
            for attr in key_attr:
                src_val = getattr(kf_lib, attr)
                if attr.startswith('handle') and 'type' not in attr:
                    src_val += Vector((key_offset, 0))

                setattr(kf, attr, src_val)

        fc.update()
    
    # redraw graph areas
    for window in bpy.context.window_manager.windows:
        screen = window.screen
        for area in screen.areas:
            if area.type == 'GRAPH_EDITOR':
                area.tag_redraw()