rigify/utils/animation.py

   1 # SPDX-FileCopyrightText: 2019-2022 Blender Foundation
   2 #
   3 # SPDX-License-Identifier: GPL-2.0-or-later
   4
   5 import bpy  # noqa
   6 import math  # noqa
   7 from mathutils import Matrix, Vector  # noqa
   8
   9 from typing import TYPE_CHECKING, Callable, Any, Collection, Iterator, Optional, Sequence
  10 from bpy.types import Action, bpy_struct, FCurve
  11
  12 import json
  13
  14 if TYPE_CHECKING:
  15     from ..rig_ui_template import PanelLayout
  16
  17
  18 rig_id = None
  19
  20
  21 ##############################################
  22 # Keyframing functions
  23 ##############################################
  24
  25 def get_keyed_frames_in_range(context, rig):
  26     action = find_action(rig)
  27     if action:
  28         frame_range = RIGIFY_OT_get_frame_range.get_range(context)
  29
  30         return sorted(get_curve_frame_set(action.fcurves, frame_range))
  31     else:
  32         return []
  33
  34
  35 def bones_in_frame(f, rig, *args):
  36     """
  37     True if one of the bones listed in args is animated at frame f
  38     :param f: the frame
  39     :param rig: the rig
  40     :param args: bone names
  41     :return:
  42     """
  43
  44     if rig.animation_data and rig.animation_data.action:
  45         fcurves = rig.animation_data.action.fcurves
  46     else:
  47         return False
  48
  49     for fc in fcurves:
  50         animated_frames = [kp.co[0] for kp in fc.keyframe_points]
  51         for bone in args:
  52             if bone in fc.data_path.split('"') and f in animated_frames:
  53                 return True
  54
  55     return False
  56
  57
  58 def overwrite_prop_animation(rig, bone, prop_name, value, frames):
  59     act = rig.animation_data.action
  60     if not act:
  61         return
  62
  63     bone_name = bone.name
  64     curve = None
  65
  66     for fcu in act.fcurves:
  67         words = fcu.data_path.split('"')
  68         if words[0] == "pose.bones[" and words[1] == bone_name and words[-2] == prop_name:
  69             curve = fcu
  70             break
  71
  72     if not curve:
  73         return
  74
  75     for kp in curve.keyframe_points:
  76         if kp.co[0] in frames:
  77             kp.co[1] = value
  78
  79
  80 ################################################################
  81 # Utilities for inserting keyframes and/or setting transforms ##
  82 ################################################################
  83
  84 SCRIPT_UTILITIES_KEYING = ['''
  85 ######################
  86 ## Keyframing tools ##
  87 ######################
  88
  89 def get_keying_flags(context):
  90     "Retrieve the general keyframing flags from user preferences."
  91     prefs = context.preferences
  92     ts = context.scene.tool_settings
  93     flags = set()
  94     # Not adding INSERTKEY_VISUAL
  95     if prefs.edit.use_keyframe_insert_needed:
  96         flags.add('INSERTKEY_NEEDED')
  97     if ts.use_keyframe_cycle_aware:
  98         flags.add('INSERTKEY_CYCLE_AWARE')
  99     return flags
 100
 101 def get_autokey_flags(context, ignore_keyingset=False):
 102     "Retrieve the Auto Keyframe flags, or None if disabled."
 103     ts = context.scene.tool_settings
 104     if ts.use_keyframe_insert_auto and (ignore_keyingset or not ts.use_keyframe_insert_keyingset):
 105         flags = get_keying_flags(context)
 106         if context.preferences.edit.use_keyframe_insert_available:
 107             flags.add('INSERTKEY_AVAILABLE')
 108         if ts.auto_keying_mode == 'REPLACE_KEYS':
 109             flags.add('INSERTKEY_REPLACE')
 110         return flags
 111     else:
 112         return None
 113
 114 def add_flags_if_set(base, new_flags):
 115     "Add more flags if base is not None."
 116     if base is None:
 117         return None
 118     else:
 119         return base | new_flags
 120
 121 def get_4d_rot_lock(bone):
 122     "Retrieve the lock status for 4D rotation."
 123     if bone.lock_rotations_4d:
 124         return [bone.lock_rotation_w, *bone.lock_rotation]
 125     else:
 126         return [all(bone.lock_rotation)] * 4
 127
 128 def keyframe_transform_properties(obj, bone_name, keyflags, *,
 129                                   ignore_locks=False, no_loc=False, no_rot=False, no_scale=False):
 130     "Keyframe transformation properties, taking flags and mode into account, and avoiding keying locked channels."
 131     bone = obj.pose.bones[bone_name]
 132
 133     def keyframe_channels(prop, locks):
 134         if ignore_locks or not all(locks):
 135             if ignore_locks or not any(locks):
 136                 bone.keyframe_insert(prop, group=bone_name, options=keyflags)
 137             else:
 138                 for i, lock in enumerate(locks):
 139                     if not lock:
 140                         bone.keyframe_insert(prop, index=i, group=bone_name, options=keyflags)
 141
 142     if not (no_loc or bone.bone.use_connect):
 143         keyframe_channels('location', bone.lock_location)
 144
 145     if not no_rot:
 146         if bone.rotation_mode == 'QUATERNION':
 147             keyframe_channels('rotation_quaternion', get_4d_rot_lock(bone))
 148         elif bone.rotation_mode == 'AXIS_ANGLE':
 149             keyframe_channels('rotation_axis_angle', get_4d_rot_lock(bone))
 150         else:
 151             keyframe_channels('rotation_euler', bone.lock_rotation)
 152
 153     if not no_scale:
 154         keyframe_channels('scale', bone.lock_scale)
 155
 156 ######################
 157 ## Constraint tools ##
 158 ######################
 159
 160 def get_constraint_target_matrix(con):
 161     target = con.target
 162     if target:
 163         if target.type == 'ARMATURE' and con.subtarget:
 164             if con.subtarget in target.pose.bones:
 165                 bone = target.pose.bones[con.subtarget]
 166                 return target.convert_space(
 167                     pose_bone=bone, matrix=bone.matrix, from_space='POSE', to_space=con.target_space)
 168         else:
 169             return target.convert_space(matrix=target.matrix_world, from_space='WORLD', to_space=con.target_space)
 170     return Matrix.Identity(4)
 171
 172 def undo_copy_scale_with_offset(obj, bone, con, old_matrix):
 173     "Undo the effects of Copy Scale with Offset constraint on a bone matrix."
 174     inf = con.influence
 175
 176     if con.mute or inf == 0 or not con.is_valid or not con.use_offset or con.use_add:
 177         return old_matrix
 178
 179     tgt_matrix = get_constraint_target_matrix(con)
 180     tgt_scale = tgt_matrix.to_scale()
 181     use = [con.use_x, con.use_y, con.use_z]
 182
 183     if con.use_make_uniform:
 184         if con.use_x and con.use_y and con.use_z:
 185             total = tgt_matrix.determinant()
 186         else:
 187             total = 1
 188             for i, use in enumerate(use):
 189                 if use:
 190                     total *= tgt_scale[i]
 191
 192         tgt_scale = [abs(total)**(1./3.)]*3
 193     else:
 194         for i, use in enumerate(use):
 195             if not use:
 196                 tgt_scale[i] = 1
 197
 198     scale_delta = [
 199         1 / (1 + (math.pow(x, con.power) - 1) * inf)
 200         for x in tgt_scale
 201     ]
 202
 203     return old_matrix @ Matrix.Diagonal([*scale_delta, 1])
 204
 205 def undo_copy_scale_constraints(obj, bone, matrix):
 206     "Undo the effects of all Copy Scale with Offset constraints on a bone matrix."
 207     for con in reversed(bone.constraints):
 208         if con.type == 'COPY_SCALE':
 209             matrix = undo_copy_scale_with_offset(obj, bone, con, matrix)
 210     return matrix
 211
 212 ###############################
 213 ## Assign and keyframe tools ##
 214 ###############################
 215
 216 def set_custom_property_value(obj, bone_name, prop, value, *, keyflags=None):
 217     "Assign the value of a custom property, and optionally keyframe it."
 218     from rna_prop_ui import rna_idprop_ui_prop_update
 219     bone = obj.pose.bones[bone_name]
 220     bone[prop] = value
 221     rna_idprop_ui_prop_update(bone, prop)
 222     if keyflags is not None:
 223         bone.keyframe_insert(rna_idprop_quote_path(prop), group=bone.name, options=keyflags)
 224
 225 def get_transform_matrix(obj, bone_name, *, space='POSE', with_constraints=True):
 226     "Retrieve the matrix of the bone before or after constraints in the given space."
 227     bone = obj.pose.bones[bone_name]
 228     if with_constraints:
 229         return obj.convert_space(pose_bone=bone, matrix=bone.matrix, from_space='POSE', to_space=space)
 230     else:
 231         return obj.convert_space(pose_bone=bone, matrix=bone.matrix_basis, from_space='LOCAL', to_space=space)
 232
 233 def get_chain_transform_matrices(obj, bone_names, **options):
 234     return [get_transform_matrix(obj, name, **options) for name in bone_names]
 235
 236 def set_transform_from_matrix(obj, bone_name, matrix, *, space='POSE', undo_copy_scale=False,
 237                               ignore_locks=False, no_loc=False, no_rot=False, no_scale=False, keyflags=None):
 238     """Apply the matrix to the transformation of the bone, taking locked channels, mode and certain
 239     constraints into account, and optionally keyframe it."""
 240     bone = obj.pose.bones[bone_name]
 241
 242     def restore_channels(prop, old_vec, locks, extra_lock):
 243         if extra_lock or (not ignore_locks and all(locks)):
 244             setattr(bone, prop, old_vec)
 245         else:
 246             if not ignore_locks and any(locks):
 247                 new_vec = Vector(getattr(bone, prop))
 248
 249                 for i, lock in enumerate(locks):
 250                     if lock:
 251                         new_vec[i] = old_vec[i]
 252
 253                 setattr(bone, prop, new_vec)
 254
 255     # Save the old values of the properties
 256     old_loc = Vector(bone.location)
 257     old_rot_euler = Vector(bone.rotation_euler)
 258     old_rot_quat = Vector(bone.rotation_quaternion)
 259     old_rot_axis = Vector(bone.rotation_axis_angle)
 260     old_scale = Vector(bone.scale)
 261
 262     # Compute and assign the local matrix
 263     if space != 'LOCAL':
 264         matrix = obj.convert_space(pose_bone=bone, matrix=matrix, from_space=space, to_space='LOCAL')
 265
 266     if undo_copy_scale:
 267         matrix = undo_copy_scale_constraints(obj, bone, matrix)
 268
 269     bone.matrix_basis = matrix
 270
 271     # Restore locked properties
 272     restore_channels('location', old_loc, bone.lock_location, no_loc or bone.bone.use_connect)
 273
 274     if bone.rotation_mode == 'QUATERNION':
 275         restore_channels('rotation_quaternion', old_rot_quat, get_4d_rot_lock(bone), no_rot)
 276         bone.rotation_axis_angle = old_rot_axis
 277         bone.rotation_euler = old_rot_euler
 278     elif bone.rotation_mode == 'AXIS_ANGLE':
 279         bone.rotation_quaternion = old_rot_quat
 280         restore_channels('rotation_axis_angle', old_rot_axis, get_4d_rot_lock(bone), no_rot)
 281         bone.rotation_euler = old_rot_euler
 282     else:
 283         bone.rotation_quaternion = old_rot_quat
 284         bone.rotation_axis_angle = old_rot_axis
 285         restore_channels('rotation_euler', old_rot_euler, bone.lock_rotation, no_rot)
 286
 287     restore_channels('scale', old_scale, bone.lock_scale, no_scale)
 288
 289     # Keyframe properties
 290     if keyflags is not None:
 291         keyframe_transform_properties(
 292             obj, bone_name, keyflags, ignore_locks=ignore_locks,
 293             no_loc=no_loc, no_rot=no_rot, no_scale=no_scale
 294         )
 295
 296 def set_chain_transforms_from_matrices(context, obj, bone_names, matrices, **options):
 297     for bone, matrix in zip(bone_names, matrices):
 298         set_transform_from_matrix(obj, bone, matrix, **options)
 299         context.view_layer.update()
 300 ''']
 301
 302 exec(SCRIPT_UTILITIES_KEYING[-1])
 303
 304 ############################################
 305 # Utilities for managing animation curves ##
 306 ############################################
 307
 308 SCRIPT_UTILITIES_CURVES = ['''
 309 ###########################
 310 ## Animation curve tools ##
 311 ###########################
 312
 313 def flatten_curve_set(curves):
 314     "Iterate over all FCurves inside a set of nested lists and dictionaries."
 315     if curves is None:
 316         pass
 317     elif isinstance(curves, bpy.types.FCurve):
 318         yield curves
 319     elif isinstance(curves, dict):
 320         for sub in curves.values():
 321             yield from flatten_curve_set(sub)
 322     else:
 323         for sub in curves:
 324             yield from flatten_curve_set(sub)
 325
 326 def flatten_curve_key_set(curves, key_range=None):
 327     "Iterate over all keys of the given fcurves in the specified range."
 328     for curve in flatten_curve_set(curves):
 329         for key in curve.keyframe_points:
 330             if key_range is None or key_range[0] <= key.co[0] <= key_range[1]:
 331                 yield key
 332
 333 def get_curve_frame_set(curves, key_range=None):
 334     "Compute a set of all time values with existing keys in the given curves and range."
 335     return set(key.co[0] for key in flatten_curve_key_set(curves, key_range))
 336
 337 def set_curve_key_interpolation(curves, ipo, key_range=None):
 338     "Assign the given interpolation value to all curve keys in range."
 339     for key in flatten_curve_key_set(curves, key_range):
 340         key.interpolation = ipo
 341
 342 def delete_curve_keys_in_range(curves, key_range=None):
 343     "Delete all keys of the given curves within the given range."
 344     for curve in flatten_curve_set(curves):
 345         points = curve.keyframe_points
 346         for i in range(len(points), 0, -1):
 347             key = points[i - 1]
 348             if key_range is None or key_range[0] <= key.co[0] <= key_range[1]:
 349                 points.remove(key, fast=True)
 350         curve.update()
 351
 352 def nla_tweak_to_scene(anim_data, frames, invert=False):
 353     "Convert a frame value or list between scene and tweaked NLA strip time."
 354     if frames is None:
 355         return None
 356     elif anim_data is None or not anim_data.use_tweak_mode:
 357         return frames
 358     elif isinstance(frames, (int, float)):
 359         return anim_data.nla_tweak_strip_time_to_scene(frames, invert=invert)
 360     else:
 361         return type(frames)(
 362             anim_data.nla_tweak_strip_time_to_scene(v, invert=invert) for v in frames
 363         )
 364
 365 def find_action(action):
 366     if isinstance(action, bpy.types.Object):
 367         action = action.animation_data
 368     if isinstance(action, bpy.types.AnimData):
 369         action = action.action
 370     if isinstance(action, bpy.types.Action):
 371         return action
 372     else:
 373         return None
 374
 375 def clean_action_empty_curves(action):
 376     "Delete completely empty curves from the given action."
 377     action = find_action(action)
 378     for curve in list(action.fcurves):
 379         if curve.is_empty:
 380             action.fcurves.remove(curve)
 381     action.update_tag()
 382
 383 TRANSFORM_PROPS_LOCATION = frozenset(['location'])
 384 TRANSFORM_PROPS_ROTATION = frozenset(['rotation_euler', 'rotation_quaternion', 'rotation_axis_angle'])
 385 TRANSFORM_PROPS_SCALE = frozenset(['scale'])
 386 TRANSFORM_PROPS_ALL = frozenset(TRANSFORM_PROPS_LOCATION | TRANSFORM_PROPS_ROTATION | TRANSFORM_PROPS_SCALE)
 387
 388 def transform_props_with_locks(lock_location, lock_rotation, lock_scale):
 389     props = set()
 390     if not lock_location:
 391         props |= TRANSFORM_PROPS_LOCATION
 392     if not lock_rotation:
 393         props |= TRANSFORM_PROPS_ROTATION
 394     if not lock_scale:
 395         props |= TRANSFORM_PROPS_SCALE
 396     return props
 397
 398 class FCurveTable(object):
 399     "Table for efficient lookup of FCurves by properties."
 400
 401     def __init__(self):
 402         self.curve_map = collections.defaultdict(dict)
 403
 404     def index_curves(self, curves):
 405         for curve in curves:
 406             index = curve.array_index
 407             if index < 0:
 408                 index = 0
 409             self.curve_map[curve.data_path][index] = curve
 410
 411     def get_prop_curves(self, ptr, prop_path):
 412         "Returns a dictionary from array index to curve for the given property, or Null."
 413         return self.curve_map.get(ptr.path_from_id(prop_path))
 414
 415     def list_all_prop_curves(self, ptr_set, path_set):
 416         "Iterates over all FCurves matching the given object(s) and properties."
 417         if isinstance(ptr_set, bpy.types.bpy_struct):
 418             ptr_set = [ptr_set]
 419         for ptr in ptr_set:
 420             for path in path_set:
 421                 curves = self.get_prop_curves(ptr, path)
 422                 if curves:
 423                     yield from curves.values()
 424
 425     def get_custom_prop_curves(self, ptr, prop):
 426         return self.get_prop_curves(ptr, rna_idprop_quote_path(prop))
 427
 428 class ActionCurveTable(FCurveTable):
 429     "Table for efficient lookup of Action FCurves by properties."
 430
 431     def __init__(self, action):
 432         super().__init__()
 433         self.action = find_action(action)
 434         if self.action:
 435             self.index_curves(self.action.fcurves)
 436
 437 class DriverCurveTable(FCurveTable):
 438     "Table for efficient lookup of Driver FCurves by properties."
 439
 440     def __init__(self, object):
 441         super().__init__()
 442         self.anim_data = object.animation_data
 443         if self.anim_data:
 444             self.index_curves(self.anim_data.drivers)
 445 ''']
 446
 447 AnyCurveSet = None | FCurve | dict | Collection
 448 flatten_curve_set: Callable[[AnyCurveSet], Iterator[FCurve]]
 449 flatten_curve_key_set: Callable[..., set[float]]
 450 get_curve_frame_set: Callable[..., set[float]]
 451 set_curve_key_interpolation: Callable[..., None]
 452 delete_curve_keys_in_range: Callable[..., None]
 453 nla_tweak_to_scene: Callable
 454 find_action: Callable[[bpy_struct], Action]
 455 clean_action_empty_curves: Callable[[bpy_struct], None]
 456 TRANSFORM_PROPS_LOCATION: frozenset[str]
 457 TRANSFORM_PROPS_ROTATION = frozenset[str]
 458 TRANSFORM_PROPS_SCALE = frozenset[str]
 459 TRANSFORM_PROPS_ALL = frozenset[str]
 460 transform_props_with_locks: Callable[[bool, bool, bool], set[str]]
 461 FCurveTable: Any
 462 ActionCurveTable: Any
 463 DriverCurveTable: Any
 464
 465 exec(SCRIPT_UTILITIES_CURVES[-1])
 466
 467 ################################################
 468 # Utilities for operators that bake keyframes ##
 469 ################################################
 470
 471 _SCRIPT_REGISTER_WM_PROPS = '''
 472 bpy.types.WindowManager.rigify_transfer_use_all_keys = bpy.props.BoolProperty(
 473     name="Bake All Keyed Frames",
 474     description="Bake on every frame that has a key for any of the bones, as opposed to just the relevant ones",
 475     default=False
 476 )
 477 bpy.types.WindowManager.rigify_transfer_use_frame_range = bpy.props.BoolProperty(
 478     name="Limit Frame Range", description="Only bake keyframes in a certain frame range", default=False
 479 )
 480 bpy.types.WindowManager.rigify_transfer_start_frame = bpy.props.IntProperty(
 481     name="Start", description="First frame to transfer", default=0, min=0
 482 )
 483 bpy.types.WindowManager.rigify_transfer_end_frame = bpy.props.IntProperty(
 484     name="End", description="Last frame to transfer", default=0, min=0
 485 )
 486 '''
 487
 488 _SCRIPT_UNREGISTER_WM_PROPS = '''
 489 del bpy.types.WindowManager.rigify_transfer_use_all_keys
 490 del bpy.types.WindowManager.rigify_transfer_use_frame_range
 491 del bpy.types.WindowManager.rigify_transfer_start_frame
 492 del bpy.types.WindowManager.rigify_transfer_end_frame
 493 '''
 494
 495 _SCRIPT_UTILITIES_BAKE_OPS = '''
 496 class RIGIFY_OT_get_frame_range(bpy.types.Operator):
 497     bl_idname = "rigify.get_frame_range" + ('_'+rig_id if rig_id else '')
 498     bl_label = "Get Frame Range"
 499     bl_description = "Set start and end frame from scene"
 500     bl_options = {'INTERNAL'}
 501
 502     def execute(self, context):
 503         scn = context.scene
 504         id_store = context.window_manager
 505         id_store.rigify_transfer_start_frame = scn.frame_start
 506         id_store.rigify_transfer_end_frame = scn.frame_end
 507         return {'FINISHED'}
 508
 509     @staticmethod
 510     def get_range(context):
 511         id_store = context.window_manager
 512         if not id_store.rigify_transfer_use_frame_range:
 513             return None
 514         else:
 515             return (id_store.rigify_transfer_start_frame, id_store.rigify_transfer_end_frame)
 516
 517     @classmethod
 518     def draw_range_ui(self, context, layout):
 519         id_store = context.window_manager
 520
 521         row = layout.row(align=True)
 522         row.prop(id_store, 'rigify_transfer_use_frame_range', icon='PREVIEW_RANGE', text='')
 523
 524         row = row.row(align=True)
 525         row.active = id_store.rigify_transfer_use_frame_range
 526         row.prop(id_store, 'rigify_transfer_start_frame')
 527         row.prop(id_store, 'rigify_transfer_end_frame')
 528         row.operator(self.bl_idname, icon='TIME', text='')
 529 '''
 530
 531 RIGIFY_OT_get_frame_range: Any
 532
 533 exec(_SCRIPT_UTILITIES_BAKE_OPS)
 534
 535 ################################################
 536 # Framework for operators that bake keyframes ##
 537 ################################################
 538
 539 SCRIPT_REGISTER_BAKE = ['RIGIFY_OT_get_frame_range']
 540
 541 SCRIPT_UTILITIES_BAKE = SCRIPT_UTILITIES_KEYING + SCRIPT_UTILITIES_CURVES + ['''
 542 ##################################
 543 # Common bake operator settings ##
 544 ##################################
 545 ''' + _SCRIPT_REGISTER_WM_PROPS + _SCRIPT_UTILITIES_BAKE_OPS + '''
 546 #######################################
 547 # Keyframe baking operator framework ##
 548 #######################################
 549
 550 class RigifyOperatorMixinBase:
 551     bl_options = {'UNDO', 'INTERNAL'}
 552
 553     def init_invoke(self, context):
 554         "Override to initialize the operator before invoke."
 555
 556     def init_execute(self, context):
 557         "Override to initialize the operator before execute."
 558
 559     def before_save_state(self, context, rig):
 560         "Override to prepare for saving state."
 561
 562     def after_save_state(self, context, rig):
 563         "Override to undo before_save_state."
 564
 565
 566 class RigifyBakeKeyframesMixin(RigifyOperatorMixinBase):
 567     """Basic framework for an operator that updates a set of keyed frames."""
 568
 569     # Utilities
 570     def nla_from_raw(self, frames):
 571         "Convert frame(s) from inner action time to scene time."
 572         return nla_tweak_to_scene(self.bake_anim, frames)
 573
 574     def nla_to_raw(self, frames):
 575         "Convert frame(s) from scene time to inner action time."
 576         return nla_tweak_to_scene(self.bake_anim, frames, invert=True)
 577
 578     def bake_get_bone(self, bone_name):
 579         "Get pose bone by name."
 580         return self.bake_rig.pose.bones[bone_name]
 581
 582     def bake_get_bones(self, bone_names):
 583         "Get multiple pose bones by name."
 584         if isinstance(bone_names, (list, set)):
 585             return [self.bake_get_bone(name) for name in bone_names]
 586         else:
 587             return self.bake_get_bone(bone_names)
 588
 589     def bake_get_all_bone_curves(self, bone_names, props):
 590         "Get a list of all curves for the specified properties of the specified bones."
 591         return list(self.bake_curve_table.list_all_prop_curves(self.bake_get_bones(bone_names), props))
 592
 593     def bake_get_all_bone_custom_prop_curves(self, bone_names, props):
 594         "Get a list of all curves for the specified custom properties of the specified bones."
 595         return self.bake_get_all_bone_curves(bone_names, [rna_idprop_quote_path(p) for p in props])
 596
 597     def bake_get_bone_prop_curves(self, bone_name, prop):
 598         "Get an index to curve dict for the specified property of the specified bone."
 599         return self.bake_curve_table.get_prop_curves(self.bake_get_bone(bone_name), prop)
 600
 601     def bake_get_bone_custom_prop_curves(self, bone_name, prop):
 602         "Get an index to curve dict for the specified custom property of the specified bone."
 603         return self.bake_curve_table.get_custom_prop_curves(self.bake_get_bone(bone_name), prop)
 604
 605     def bake_add_curve_frames(self, curves):
 606         "Register frames keyed in the specified curves for baking."
 607         self.bake_frames_raw |= get_curve_frame_set(curves, self.bake_frame_range_raw)
 608
 609     def bake_add_bone_frames(self, bone_names, props):
 610         "Register frames keyed for the specified properties of the specified bones for baking."
 611         curves = self.bake_get_all_bone_curves(bone_names, props)
 612         self.bake_add_curve_frames(curves)
 613         return curves
 614
 615     def bake_replace_custom_prop_keys_constant(self, bone, prop, new_value):
 616         "If the property is keyframed, delete keys in bake range and re-key as Constant."
 617         prop_curves = self.bake_get_bone_custom_prop_curves(bone, prop)
 618
 619         if prop_curves and 0 in prop_curves:
 620             range_raw = self.nla_to_raw(self.get_bake_range())
 621             delete_curve_keys_in_range(prop_curves, range_raw)
 622             set_custom_property_value(self.bake_rig, bone, prop, new_value, keyflags={'INSERTKEY_AVAILABLE'})
 623             set_curve_key_interpolation(prop_curves, 'CONSTANT', range_raw)
 624
 625     # Default behavior implementation
 626     def bake_init(self, context):
 627         self.bake_rig = context.active_object
 628         self.bake_anim = self.bake_rig.animation_data
 629         self.bake_frame_range = RIGIFY_OT_get_frame_range.get_range(context)
 630         self.bake_frame_range_raw = self.nla_to_raw(self.bake_frame_range)
 631         self.bake_curve_table = ActionCurveTable(self.bake_rig)
 632         self.bake_current_frame = context.scene.frame_current
 633         self.bake_frames_raw = set()
 634         self.bake_state = dict()
 635
 636         self.keyflags = get_keying_flags(context)
 637         self.keyflags_switch = None
 638
 639         if context.window_manager.rigify_transfer_use_all_keys:
 640             self.bake_add_curve_frames(self.bake_curve_table.curve_map)
 641
 642     def bake_add_frames_done(self):
 643         "Computes and sets the final set of frames to bake."
 644         frames = self.nla_from_raw(self.bake_frames_raw)
 645         self.bake_frames = sorted(set(map(round, frames)))
 646
 647     def is_bake_empty(self):
 648         return len(self.bake_frames_raw) == 0
 649
 650     def report_bake_empty(self):
 651         self.bake_add_frames_done()
 652         if self.is_bake_empty():
 653             self.report({'WARNING'}, 'No keys to bake.')
 654             return True
 655         return False
 656
 657     def get_bake_range(self):
 658         "Returns the frame range that is being baked."
 659         if self.bake_frame_range:
 660             return self.bake_frame_range
 661         else:
 662             frames = self.bake_frames
 663             return (frames[0], frames[-1])
 664
 665     def get_bake_range_pair(self):
 666         "Returns the frame range that is being baked, both in scene and action time."
 667         range = self.get_bake_range()
 668         return range, self.nla_to_raw(range)
 669
 670     def bake_save_state(self, context):
 671         "Scans frames and collects data for baking before changing anything."
 672         rig = self.bake_rig
 673         scene = context.scene
 674         saved_state = self.bake_state
 675
 676         try:
 677             self.before_save_state(context, rig)
 678
 679             for frame in self.bake_frames:
 680                 scene.frame_set(frame)
 681                 saved_state[frame] = self.save_frame_state(context, rig)
 682
 683         finally:
 684             self.after_save_state(context, rig)
 685
 686     def bake_clean_curves_in_range(self, context, curves):
 687         "Deletes all keys from the given curves in the bake range."
 688         range, range_raw = self.get_bake_range_pair()
 689
 690         context.scene.frame_set(range[0])
 691         delete_curve_keys_in_range(curves, range_raw)
 692
 693         return range, range_raw
 694
 695     def bake_apply_state(self, context):
 696         "Scans frames and applies the baking operation."
 697         rig = self.bake_rig
 698         scene = context.scene
 699         saved_state = self.bake_state
 700
 701         for frame in self.bake_frames:
 702             scene.frame_set(frame)
 703             self.apply_frame_state(context, rig, saved_state.get(frame))
 704
 705         clean_action_empty_curves(self.bake_rig)
 706         scene.frame_set(self.bake_current_frame)
 707
 708     @staticmethod
 709     def draw_common_bake_ui(context, layout):
 710         layout.prop(context.window_manager, 'rigify_transfer_use_all_keys')
 711
 712         RIGIFY_OT_get_frame_range.draw_range_ui(context, layout)
 713
 714     @classmethod
 715     def poll(cls, context):
 716         return find_action(context.active_object) is not None
 717
 718     def execute_scan_curves(self, context, obj):
 719         "Override to register frames to be baked, and return curves that should be cleared."
 720         raise NotImplementedError()
 721
 722     def execute_before_apply(self, context, obj, range, range_raw):
 723         "Override to execute code one time before the bake apply frame scan."
 724         pass
 725
 726     def execute(self, context):
 727         self.init_execute(context)
 728         self.bake_init(context)
 729
 730         curves = self.execute_scan_curves(context, self.bake_rig)
 731
 732         if self.report_bake_empty():
 733             return {'CANCELLED'}
 734
 735         try:
 736             self.bake_save_state(context)
 737
 738             range, range_raw = self.bake_clean_curves_in_range(context, curves)
 739
 740             self.execute_before_apply(context, self.bake_rig, range, range_raw)
 741
 742             self.bake_apply_state(context)
 743
 744         except Exception as e:
 745             traceback.print_exc()
 746             self.report({'ERROR'}, 'Exception: ' + str(e))
 747
 748         return {'FINISHED'}
 749
 750     def invoke(self, context, event):
 751         self.init_invoke(context)
 752
 753         if hasattr(self, 'draw'):
 754             return context.window_manager.invoke_props_dialog(self)
 755         else:
 756             return context.window_manager.invoke_confirm(self, event)
 757
 758
 759 class RigifySingleUpdateMixin(RigifyOperatorMixinBase):
 760     """Basic framework for an operator that updates only the current frame."""
 761
 762     def execute(self, context):
 763         self.init_execute(context)
 764         obj = context.active_object
 765         self.keyflags = get_autokey_flags(context, ignore_keyingset=True)
 766         self.keyflags_switch = add_flags_if_set(self.keyflags, {'INSERTKEY_AVAILABLE'})
 767
 768         try:
 769             try:
 770                 self.before_save_state(context, obj)
 771                 state = self.save_frame_state(context, obj)
 772             finally:
 773                 self.after_save_state(context, obj)
 774
 775             self.apply_frame_state(context, obj, state)
 776
 777         except Exception as e:
 778             traceback.print_exc()
 779             self.report({'ERROR'}, 'Exception: ' + str(e))
 780
 781         return {'FINISHED'}
 782
 783     def invoke(self, context, event):
 784         self.init_invoke(context)
 785
 786         if hasattr(self, 'draw'):
 787             return context.window_manager.invoke_props_popup(self, event)
 788         else:
 789             return self.execute(context)
 790 ''']
 791
 792 RigifyOperatorMixinBase: Any
 793 RigifyBakeKeyframesMixin: Any
 794 RigifySingleUpdateMixin: Any
 795
 796 exec(SCRIPT_UTILITIES_BAKE[-1])
 797
 798 #####################################
 799 # Generic Clear Keyframes operator ##
 800 #####################################
 801
 802 SCRIPT_REGISTER_OP_CLEAR_KEYS = ['POSE_OT_rigify_clear_keyframes']
 803
 804 SCRIPT_UTILITIES_OP_CLEAR_KEYS = ['''
 805 #############################
 806 ## Generic Clear Keyframes ##
 807 #############################
 808
 809 class POSE_OT_rigify_clear_keyframes(bpy.types.Operator):
 810     bl_idname = "pose.rigify_clear_keyframes_" + rig_id
 811     bl_label = "Clear Keyframes And Transformation"
 812     bl_options = {'UNDO', 'INTERNAL'}
 813     bl_description = "Remove all keyframes for the relevant bones and reset transformation"
 814
 815     bones: StringProperty(name="Bone List")
 816
 817     @classmethod
 818     def poll(cls, context):
 819         return find_action(context.active_object) is not None
 820
 821     def invoke(self, context, event):
 822         return context.window_manager.invoke_confirm(self, event)
 823
 824     def execute(self, context):
 825         obj = context.active_object
 826         bone_list = [ obj.pose.bones[name] for name in json.loads(self.bones) ]
 827
 828         curve_table = ActionCurveTable(context.active_object)
 829         curves = list(curve_table.list_all_prop_curves(bone_list, TRANSFORM_PROPS_ALL))
 830
 831         key_range = RIGIFY_OT_get_frame_range.get_range(context)
 832         range_raw = nla_tweak_to_scene(obj.animation_data, key_range, invert=True)
 833         delete_curve_keys_in_range(curves, range_raw)
 834
 835         for bone in bone_list:
 836             bone.location = bone.rotation_euler = (0,0,0)
 837             bone.rotation_quaternion = (1,0,0,0)
 838             bone.rotation_axis_angle = (0,0,1,0)
 839             bone.scale = (1,1,1)
 840
 841         clean_action_empty_curves(obj)
 842         obj.update_tag(refresh={'TIME'})
 843         return {'FINISHED'}
 844 ''']
 845
 846
 847 def add_clear_keyframes_button(panel: 'PanelLayout', *,
 848                                bones: Sequence[str] = (), text=''):
 849     panel.use_bake_settings()
 850     panel.script.add_utilities(SCRIPT_UTILITIES_OP_CLEAR_KEYS)
 851     panel.script.register_classes(SCRIPT_REGISTER_OP_CLEAR_KEYS)
 852
 853     op_props = {'bones': json.dumps(bones)}
 854
 855     panel.operator('pose.rigify_clear_keyframes_{rig_id}', text=text, icon='CANCEL',
 856                    properties=op_props)
 857
 858
 859 ###################################
 860 # Generic Snap FK to IK operator ##
 861 ###################################
 862
 863 SCRIPT_REGISTER_OP_SNAP = ['POSE_OT_rigify_generic_snap', 'POSE_OT_rigify_generic_snap_bake']
 864
 865 SCRIPT_UTILITIES_OP_SNAP = ['''
 866 #############################
 867 ## Generic Snap (FK to IK) ##
 868 #############################
 869
 870 class RigifyGenericSnapBase:
 871     input_bones:   StringProperty(name="Input Chain")
 872     output_bones:  StringProperty(name="Output Chain")
 873     ctrl_bones:    StringProperty(name="Input Controls")
 874
 875     tooltip:         StringProperty(name="Tooltip", default="FK to IK")
 876     locks:           bpy.props.BoolVectorProperty(name="Locked", size=3, default=[False,False,False])
 877     undo_copy_scale: bpy.props.BoolProperty(name="Undo Copy Scale", default=False)
 878
 879     def init_execute(self, context):
 880         self.input_bone_list = json.loads(self.input_bones)
 881         self.output_bone_list = json.loads(self.output_bones)
 882         self.ctrl_bone_list = json.loads(self.ctrl_bones)
 883
 884     def save_frame_state(self, context, obj):
 885         return get_chain_transform_matrices(obj, self.input_bone_list)
 886
 887     def apply_frame_state(self, context, obj, matrices):
 888         set_chain_transforms_from_matrices(
 889             context, obj, self.output_bone_list, matrices,
 890             undo_copy_scale=self.undo_copy_scale, keyflags=self.keyflags,
 891             no_loc=self.locks[0], no_rot=self.locks[1], no_scale=self.locks[2],
 892         )
 893
 894 class POSE_OT_rigify_generic_snap(RigifyGenericSnapBase, RigifySingleUpdateMixin, bpy.types.Operator):
 895     bl_idname = "pose.rigify_generic_snap_" + rig_id
 896     bl_label = "Snap Bones"
 897     bl_description = "Snap on the current frame"
 898
 899     @classmethod
 900     def description(cls, context, props):
 901         return "Snap " + props.tooltip + " on the current frame"
 902
 903 class POSE_OT_rigify_generic_snap_bake(RigifyGenericSnapBase, RigifyBakeKeyframesMixin, bpy.types.Operator):
 904     bl_idname = "pose.rigify_generic_snap_bake_" + rig_id
 905     bl_label = "Apply Snap To Keyframes"
 906     bl_description = "Apply snap to keyframes"
 907
 908     @classmethod
 909     def description(cls, context, props):
 910         return "Apply snap " + props.tooltip + " to keyframes"
 911
 912     def execute_scan_curves(self, context, obj):
 913         props = transform_props_with_locks(*self.locks)
 914         self.bake_add_bone_frames(self.ctrl_bone_list, TRANSFORM_PROPS_ALL)
 915         return self.bake_get_all_bone_curves(self.output_bone_list, props)
 916 ''']
 917
 918
 919 def add_fk_ik_snap_buttons(panel: 'PanelLayout', op_single: str, op_bake: str, *,
 920                            label, rig_name='', properties: dict[str, Any],
 921                            clear_bones: Optional[list[str]] = None,
 922                            compact: Optional[bool] = None):
 923     assert label and properties
 924
 925     if rig_name:
 926         label += ' (%s)' % rig_name
 927
 928     if compact or not clear_bones:
 929         row = panel.row(align=True)
 930         row.operator(op_single, text=label, icon='SNAP_ON', properties=properties)
 931         row.operator(op_bake, text='', icon='ACTION_TWEAK', properties=properties)
 932
 933         if clear_bones:
 934             add_clear_keyframes_button(row, bones=clear_bones)
 935     else:
 936         col = panel.column(align=True)
 937         col.operator(op_single, text=label, icon='SNAP_ON', properties=properties)
 938         row = col.row(align=True)
 939         row.operator(op_bake, text='Action', icon='ACTION_TWEAK', properties=properties)
 940         add_clear_keyframes_button(row, bones=clear_bones, text='Clear')
 941
 942
 943 def add_generic_snap(panel: 'PanelLayout', *,
 944                      output_bones: Sequence[str] = (), input_bones: Sequence[str] = (),
 945                      input_ctrl_bones: Sequence[str] = (), label='Snap',
 946                      rig_name='', undo_copy_scale=False, compact: Optional[bool] = None,
 947                      clear=True, locks: Optional[Sequence[bool]] = None,
 948                      tooltip: Optional[str] = None):
 949     panel.use_bake_settings()
 950     panel.script.add_utilities(SCRIPT_UTILITIES_OP_SNAP)
 951     panel.script.register_classes(SCRIPT_REGISTER_OP_SNAP)
 952
 953     op_props = {
 954         'output_bones': json.dumps(output_bones),
 955         'input_bones': json.dumps(input_bones),
 956         'ctrl_bones': json.dumps(input_ctrl_bones or input_bones),
 957     }
 958
 959     if undo_copy_scale:
 960         op_props['undo_copy_scale'] = undo_copy_scale
 961     if locks is not None:
 962         op_props['locks'] = tuple(locks[0:3])
 963     if tooltip is not None:
 964         op_props['tooltip'] = tooltip
 965
 966     clear_bones = output_bones if clear else None
 967
 968     add_fk_ik_snap_buttons(
 969         panel, 'pose.rigify_generic_snap_{rig_id}', 'pose.rigify_generic_snap_bake_{rig_id}',
 970         label=label, rig_name=rig_name, properties=op_props, clear_bones=clear_bones, compact=compact,
 971     )
 972
 973
 974 def add_generic_snap_fk_to_ik(panel: 'PanelLayout', *,
 975                               fk_bones: Sequence[str] = (), ik_bones: Sequence[str] = (),
 976                               ik_ctrl_bones: Sequence[str] = (), label='FK->IK',
 977                               rig_name='', undo_copy_scale=False,
 978                               compact: Optional[bool] = None, clear=True):
 979     add_generic_snap(
 980         panel, output_bones=fk_bones, input_bones=ik_bones, input_ctrl_bones=ik_ctrl_bones,
 981         label=label, rig_name=rig_name, undo_copy_scale=undo_copy_scale, compact=compact, clear=clear
 982     )
 983
 984
 985 ###############################
 986 # Module register/unregister ##
 987 ###############################
 988
 989 def register():
 990     from bpy.utils import register_class
 991
 992     exec(_SCRIPT_REGISTER_WM_PROPS)
 993
 994     register_class(RIGIFY_OT_get_frame_range)
 995
 996
 997 def unregister():
 998     from bpy.utils import unregister_class
 999
1000     exec(_SCRIPT_UNREGISTER_WM_PROPS)
1001
1002     unregister_class(RIGIFY_OT_get_frame_range)