pose_library/pose_creation.py

   1 # SPDX-License-Identifier: GPL-2.0-or-later
   2
   3 """
   4 Pose Library - creation functions.
   5 """
   6
   7 import dataclasses
   8 import functools
   9 import re
  10
  11 from typing import Optional, FrozenSet, Set, Union, Iterable, cast
  12
  13 if "functions" not in locals():
  14     from . import asset_browser, functions
  15 else:
  16     import importlib
  17
  18     asset_browser = importlib.reload(asset_browser)
  19     functions = importlib.reload(functions)
  20
  21 import bpy
  22 from bpy.types import (
  23     Action,
  24     Bone,
  25     Context,
  26     FCurve,
  27     Keyframe,
  28 )
  29
  30 FCurveValue = Union[float, int]
  31
  32 pose_bone_re = re.compile(r'pose.bones\["([^"]+)"\]')
  33 """RegExp for matching FCurve data paths."""
  34
  35
  36 @dataclasses.dataclass(unsafe_hash=True, frozen=True)
  37 class PoseCreationParams:
  38     armature_ob: bpy.types.Object
  39     src_action: Optional[Action]
  40     src_frame_nr: float
  41     bone_names: FrozenSet[str]
  42     new_asset_name: str
  43
  44
  45 class UnresolvablePathError(ValueError):
  46     """Raised when a data_path cannot be resolved to a current value."""
  47
  48
  49 @dataclasses.dataclass(unsafe_hash=True)
  50 class PoseActionCreator:
  51     """Create an Action that's suitable for marking as Asset.
  52
  53     Does not mark as asset yet, nor does it add asset metadata.
  54     """
  55
  56     params: PoseCreationParams
  57
  58     # These were taken from Blender's Action baking code in `anim_utils.py`.
  59     # Items are (name, array_length) tuples.
  60     _bbone_props = [
  61         ("bbone_curveinx", None),
  62         ("bbone_curveoutx", None),
  63         ("bbone_curveinz", None),
  64         ("bbone_curveoutz", None),
  65         ("bbone_rollin", None),
  66         ("bbone_rollout", None),
  67         ("bbone_scalein", 3),
  68         ("bbone_scaleout", 3),
  69         ("bbone_easein", None),
  70         ("bbone_easeout", None),
  71     ]
  72
  73     def create(self) -> Optional[Action]:
  74         """Create a single-frame Action containing only the given bones, or None if no anim data was found."""
  75
  76         try:
  77             dst_action = self._create_new_action()
  78             self._store_pose(dst_action)
  79         finally:
  80             # Prevent next instantiations of this class from reusing pointers to
  81             # bones. They may not be valid by then any more.
  82             self._find_bone.cache_clear()
  83
  84         if len(dst_action.fcurves) == 0:
  85             bpy.data.actions.remove(dst_action)
  86             return None
  87
  88         return dst_action
  89
  90     def _create_new_action(self) -> Action:
  91         dst_action = bpy.data.actions.new(self.params.new_asset_name)
  92         if self.params.src_action:
  93             dst_action.id_root = self.params.src_action.id_root
  94         dst_action.user_clear()  # actions.new() sets users=1, but marking as asset also increments user count.
  95         return dst_action
  96
  97     def _store_pose(self, dst_action: Action) -> None:
  98         """Store the current pose into the given action."""
  99         self._store_bone_pose_parameters(dst_action)
 100         self._store_animated_parameters(dst_action)
 101
 102     def _store_bone_pose_parameters(self, dst_action: Action) -> None:
 103         """Store loc/rot/scale/bbone values in the Action."""
 104
 105         for bone_name in sorted(self.params.bone_names):
 106             self._store_location(dst_action, bone_name)
 107             self._store_rotation(dst_action, bone_name)
 108             self._store_scale(dst_action, bone_name)
 109             self._store_bbone(dst_action, bone_name)
 110
 111     def _store_animated_parameters(self, dst_action: Action) -> None:
 112         """Store the current value of any animated bone properties."""
 113         if self.params.src_action is None:
 114             return
 115
 116         armature_ob = self.params.armature_ob
 117         for fcurve in self.params.src_action.fcurves:
 118             match = pose_bone_re.match(fcurve.data_path)
 119             if not match:
 120                 # Not animating a bone property.
 121                 continue
 122
 123             bone_name = match.group(1)
 124             if bone_name not in self.params.bone_names:
 125                 # Bone is not our export set.
 126                 continue
 127
 128             if dst_action.fcurves.find(fcurve.data_path, index=fcurve.array_index):
 129                 # This property is already handled by a previous _store_xxx() call.
 130                 continue
 131
 132             # Only include in the pose if there is a key on this frame.
 133             if not self._has_key_on_frame(fcurve):
 134                 continue
 135
 136             try:
 137                 value = self._current_value(armature_ob, fcurve.data_path, fcurve.array_index)
 138             except UnresolvablePathError:
 139                 # A once-animated property no longer exists.
 140                 continue
 141
 142             dst_fcurve = dst_action.fcurves.new(fcurve.data_path, index=fcurve.array_index, action_group=bone_name)
 143             dst_fcurve.keyframe_points.insert(self.params.src_frame_nr, value=value)
 144             dst_fcurve.update()
 145
 146     def _store_location(self, dst_action: Action, bone_name: str) -> None:
 147         """Store bone location."""
 148         self._store_bone_array(dst_action, bone_name, "location", 3)
 149
 150     def _store_rotation(self, dst_action: Action, bone_name: str) -> None:
 151         """Store bone rotation given current rotation mode."""
 152         bone = self._find_bone(bone_name)
 153         if bone.rotation_mode == "QUATERNION":
 154             self._store_bone_array(dst_action, bone_name, "rotation_quaternion", 4)
 155         elif bone.rotation_mode == "AXIS_ANGLE":
 156             self._store_bone_array(dst_action, bone_name, "rotation_axis_angle", 4)
 157         else:
 158             self._store_bone_array(dst_action, bone_name, "rotation_euler", 3)
 159
 160     def _store_scale(self, dst_action: Action, bone_name: str) -> None:
 161         """Store bone scale."""
 162         self._store_bone_array(dst_action, bone_name, "scale", 3)
 163
 164     def _store_bbone(self, dst_action: Action, bone_name: str) -> None:
 165         """Store bendy-bone parameters."""
 166         for prop_name, array_length in self._bbone_props:
 167             if array_length:
 168                 self._store_bone_array(dst_action, bone_name, prop_name, array_length)
 169             else:
 170                 self._store_bone_property(dst_action, bone_name, prop_name)
 171
 172     def _store_bone_array(self, dst_action: Action, bone_name: str, property_name: str, array_length: int) -> None:
 173         """Store all elements of an array property."""
 174         for array_index in range(array_length):
 175             self._store_bone_property(dst_action, bone_name, property_name, array_index)
 176
 177     def _store_bone_property(
 178         self,
 179         dst_action: Action,
 180         bone_name: str,
 181         property_path: str,
 182         array_index: int = -1,
 183     ) -> None:
 184         """Store the current value of a single bone property."""
 185
 186         bone = self._find_bone(bone_name)
 187         value = self._current_value(bone, property_path, array_index)
 188
 189         # Get the full 'pose.bones["bone_name"].blablabla' path suitable for FCurves.
 190         rna_path = bone.path_from_id(property_path)
 191
 192         fcurve: Optional[FCurve] = dst_action.fcurves.find(rna_path, index=array_index)
 193         if fcurve is None:
 194             fcurve = dst_action.fcurves.new(rna_path, index=array_index, action_group=bone_name)
 195
 196         fcurve.keyframe_points.insert(self.params.src_frame_nr, value=value)
 197         fcurve.update()
 198
 199     @classmethod
 200     def _current_value(cls, datablock: bpy.types.ID, data_path: str, array_index: int) -> FCurveValue:
 201         """Resolve an RNA path + array index to an actual value."""
 202         value_or_array = cls._path_resolve(datablock, data_path)
 203
 204         # Both indices -1 and 0 are used for non-array properties.
 205         # -1 cannot be used in arrays, whereas 0 can be used in both arrays and non-arrays.
 206
 207         if array_index == -1:
 208             return cast(FCurveValue, value_or_array)
 209
 210         if array_index == 0:
 211             value_or_array = cls._path_resolve(datablock, data_path)
 212             try:
 213                 # MyPy doesn't understand this try/except is to determine the type.
 214                 value = value_or_array[array_index]  # type: ignore
 215             except TypeError:
 216                 # Not an array after all.
 217                 return cast(FCurveValue, value_or_array)
 218             return cast(FCurveValue, value)
 219
 220         # MyPy doesn't understand that array_index>0 implies this is indexable.
 221         return cast(FCurveValue, value_or_array[array_index])  # type: ignore
 222
 223     @staticmethod
 224     def _path_resolve(datablock: bpy.types.ID, data_path: str) -> Union[FCurveValue, Iterable[FCurveValue]]:
 225         """Wrapper for datablock.path_resolve(data_path).
 226
 227         Raise UnresolvablePathError when the path cannot be resolved.
 228         This is easier to deal with upstream than the generic ValueError raised
 229         by Blender.
 230         """
 231         try:
 232             return datablock.path_resolve(data_path)  # type: ignore
 233         except ValueError as ex:
 234             raise UnresolvablePathError(str(ex)) from ex
 235
 236     @functools.lru_cache(maxsize=1024)
 237     def _find_bone(self, bone_name: str) -> Bone:
 238         """Find a bone by name.
 239
 240         Assumes the named bone exists, as the bones this class handles comes
 241         from the user's selection, and you can't select a non-existent bone.
 242         """
 243
 244         bone: Bone = self.params.armature_ob.pose.bones[bone_name]
 245         return bone
 246
 247     def _has_key_on_frame(self, fcurve: FCurve) -> bool:
 248         """Return True iff the FCurve has a key on the source frame."""
 249
 250         points = fcurve.keyframe_points
 251         if not points:
 252             return False
 253
 254         frame_to_find = self.params.src_frame_nr
 255         margin = 0.001
 256         high = len(points) - 1
 257         low = 0
 258         while low <= high:
 259             mid = (high + low) // 2
 260             diff = points[mid].co.x - frame_to_find
 261             if abs(diff) < margin:
 262                 return True
 263             if diff < 0:
 264                 # Frame to find is bigger than the current middle.
 265                 low = mid + 1
 266             else:
 267                 # Frame to find is smaller than the current middle
 268                 high = mid - 1
 269         return False
 270
 271
 272 def create_pose_asset(
 273     params: PoseCreationParams,
 274 ) -> Optional[Action]:
 275     """Create a single-frame Action containing only the pose of the given bones.
 276
 277     DOES mark as asset, DOES NOT configure asset metadata.
 278     """
 279
 280     creator = PoseActionCreator(params)
 281     pose_action = creator.create()
 282     if pose_action is None:
 283         return None
 284
 285     pose_action.asset_mark()
 286     pose_action.asset_generate_preview()
 287     return pose_action
 288
 289
 290 def create_pose_asset_from_context(context: Context, new_asset_name: str) -> Optional[Action]:
 291     """Create Action asset from active object & selected bones."""
 292
 293     bones = context.selected_pose_bones_from_active_object
 294     bone_names = {bone.name for bone in bones}
 295
 296     params = PoseCreationParams(
 297         context.object,
 298         getattr(context.object.animation_data, "action", None),
 299         context.scene.frame_current,
 300         frozenset(bone_names),
 301         new_asset_name,
 302     )
 303
 304     return create_pose_asset(params)
 305
 306
 307 def copy_fcurves(
 308     dst_action: Action,
 309     src_action: Action,
 310     src_frame_nr: float,
 311     bone_names: Set[str],
 312 ) -> int:
 313     """Copy FCurves, returning number of curves copied."""
 314     num_fcurves_copied = 0
 315     for fcurve in src_action.fcurves:
 316         match = pose_bone_re.match(fcurve.data_path)
 317         if not match:
 318             continue
 319
 320         bone_name = match.group(1)
 321         if bone_name not in bone_names:
 322             continue
 323
 324         # Check if there is a keyframe on this frame.
 325         keyframe = find_keyframe(fcurve, src_frame_nr)
 326         if keyframe is None:
 327             continue
 328         create_single_key_fcurve(dst_action, fcurve, keyframe)
 329         num_fcurves_copied += 1
 330     return num_fcurves_copied
 331
 332
 333 def create_single_key_fcurve(dst_action: Action, src_fcurve: FCurve, src_keyframe: Keyframe) -> FCurve:
 334     """Create a copy of the source FCurve, but only for the given keyframe.
 335
 336     Returns a new FCurve with just one keyframe.
 337     """
 338
 339     dst_fcurve = copy_fcurve_without_keys(dst_action, src_fcurve)
 340     copy_keyframe(dst_fcurve, src_keyframe)
 341     return dst_fcurve
 342
 343
 344 def copy_fcurve_without_keys(dst_action: Action, src_fcurve: FCurve) -> FCurve:
 345     """Create a new FCurve and copy some properties."""
 346
 347     src_group_name = src_fcurve.group.name if src_fcurve.group else ""
 348     dst_fcurve = dst_action.fcurves.new(src_fcurve.data_path, index=src_fcurve.array_index, action_group=src_group_name)
 349     for propname in {"auto_smoothing", "color", "color_mode", "extrapolation"}:
 350         setattr(dst_fcurve, propname, getattr(src_fcurve, propname))
 351     return dst_fcurve
 352
 353
 354 def copy_keyframe(dst_fcurve: FCurve, src_keyframe: Keyframe) -> Keyframe:
 355     """Copy a keyframe from one FCurve to the other."""
 356
 357     dst_keyframe = dst_fcurve.keyframe_points.insert(
 358         src_keyframe.co.x, src_keyframe.co.y, options={'FAST'}, keyframe_type=src_keyframe.type
 359     )
 360
 361     for propname in {
 362         "amplitude",
 363         "back",
 364         "easing",
 365         "handle_left",
 366         "handle_left_type",
 367         "handle_right",
 368         "handle_right_type",
 369         "interpolation",
 370         "period",
 371     }:
 372         setattr(dst_keyframe, propname, getattr(src_keyframe, propname))
 373     dst_fcurve.update()
 374     return dst_keyframe
 375
 376
 377 def find_keyframe(fcurve: FCurve, frame: float) -> Optional[Keyframe]:
 378     # Binary search adapted from https://pythonguides.com/python-binary-search/
 379     keyframes = fcurve.keyframe_points
 380     low = 0
 381     high = len(keyframes) - 1
 382     mid = 0
 383
 384     # Accept any keyframe that's within 'epsilon' of the requested frame.
 385     # This should account for rounding errors and the likes.
 386     epsilon = 1e-4
 387     frame_lowerbound = frame - epsilon
 388     frame_upperbound = frame + epsilon
 389     while low <= high:
 390         mid = (high + low) // 2
 391         keyframe = keyframes[mid]
 392         if keyframe.co.x < frame_lowerbound:
 393             low = mid + 1
 394         elif keyframe.co.x > frame_upperbound:
 395             high = mid - 1
 396         else:
 397             return keyframe
 398     return None
 399
 400
 401 def assign_from_asset_browser(asset: Action, asset_browser_area: bpy.types.Area) -> None:
 402     """Assign some things from the asset browser to the asset.
 403
 404     This sets the current catalog ID, and in the future could include tags
 405     from the active dynamic catalog, etc.
 406     """
 407
 408     cat_id = asset_browser.active_catalog_id(asset_browser_area)
 409     asset.asset_data.catalog_id = cat_id