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         self._store_parameters_from_callback(dst_action)
 102
 103     def _store_bone_pose_parameters(self, dst_action: Action) -> None:
 104         """Store loc/rot/scale/bbone values in the Action."""
 105
 106         for bone_name in sorted(self.params.bone_names):
 107             self._store_location(dst_action, bone_name)
 108             self._store_rotation(dst_action, bone_name)
 109             self._store_scale(dst_action, bone_name)
 110             self._store_bbone(dst_action, bone_name)
 111
 112     def _store_animated_parameters(self, dst_action: Action) -> None:
 113         """Store the current value of any animated bone properties."""
 114         if self.params.src_action is None:
 115             return
 116
 117         armature_ob = self.params.armature_ob
 118         for fcurve in self.params.src_action.fcurves:
 119             match = pose_bone_re.match(fcurve.data_path)
 120             if not match:
 121                 # Not animating a bone property.
 122                 continue
 123
 124             bone_name = match.group(1)
 125             if bone_name not in self.params.bone_names:
 126                 # Bone is not our export set.
 127                 continue
 128
 129             if dst_action.fcurves.find(fcurve.data_path, index=fcurve.array_index):
 130                 # This property is already handled by a previous _store_xxx() call.
 131                 continue
 132
 133             # Only include in the pose if there is a key on this frame.
 134             if not self._has_key_on_frame(fcurve):
 135                 continue
 136
 137             try:
 138                 value = self._current_value(armature_ob, fcurve.data_path, fcurve.array_index)
 139             except UnresolvablePathError:
 140                 # A once-animated property no longer exists.
 141                 continue
 142
 143             dst_fcurve = dst_action.fcurves.new(
 144                 fcurve.data_path, index=fcurve.array_index, action_group=bone_name
 145             )
 146             dst_fcurve.keyframe_points.insert(self.params.src_frame_nr, value=value)
 147             dst_fcurve.update()
 148
 149     def _store_parameters_from_callback(self, dst_action: Action) -> None:
 150         """Store extra parameters in the pose based on arbitrary callbacks.
 151
 152         Not implemented yet, needs a proper design & some user stories.
 153         """
 154         pass
 155
 156     def _store_location(self, dst_action: Action, bone_name: str) -> None:
 157         """Store bone location."""
 158         self._store_bone_array(dst_action, bone_name, "location", 3)
 159
 160     def _store_rotation(self, dst_action: Action, bone_name: str) -> None:
 161         """Store bone rotation given current rotation mode."""
 162         bone = self._find_bone(bone_name)
 163         if bone.rotation_mode == "QUATERNION":
 164             self._store_bone_array(dst_action, bone_name, "rotation_quaternion", 4)
 165         elif bone.rotation_mode == "AXIS_ANGLE":
 166             self._store_bone_array(dst_action, bone_name, "rotation_axis_angle", 4)
 167         else:
 168             self._store_bone_array(dst_action, bone_name, "rotation_euler", 3)
 169
 170     def _store_scale(self, dst_action: Action, bone_name: str) -> None:
 171         """Store bone scale."""
 172         self._store_bone_array(dst_action, bone_name, "scale", 3)
 173
 174     def _store_bbone(self, dst_action: Action, bone_name: str) -> None:
 175         """Store bendy-bone parameters."""
 176         for prop_name, array_length in self._bbone_props:
 177             if array_length:
 178                 self._store_bone_array(dst_action, bone_name, prop_name, array_length)
 179             else:
 180                 self._store_bone_property(dst_action, bone_name, prop_name)
 181
 182     def _store_bone_array(
 183         self, dst_action: Action, bone_name: str, property_name: str, array_length: int
 184     ) -> None:
 185         """Store all elements of an array property."""
 186         for array_index in range(array_length):
 187             self._store_bone_property(dst_action, bone_name, property_name, array_index)
 188
 189     def _store_bone_property(
 190         self,
 191         dst_action: Action,
 192         bone_name: str,
 193         property_path: str,
 194         array_index: int = -1,
 195     ) -> None:
 196         """Store the current value of a single bone property."""
 197
 198         bone = self._find_bone(bone_name)
 199         value = self._current_value(bone, property_path, array_index)
 200
 201         # Get the full 'pose.bones["bone_name"].blablabla' path suitable for FCurves.
 202         rna_path = bone.path_from_id(property_path)
 203
 204         fcurve: Optional[FCurve] = dst_action.fcurves.find(rna_path, index=array_index)
 205         if fcurve is None:
 206             fcurve = dst_action.fcurves.new(rna_path, index=array_index, action_group=bone_name)
 207
 208         fcurve.keyframe_points.insert(self.params.src_frame_nr, value=value)
 209         fcurve.update()
 210
 211     @classmethod
 212     def _current_value(
 213         cls, datablock: bpy.types.ID, data_path: str, array_index: int
 214     ) -> FCurveValue:
 215         """Resolve an RNA path + array index to an actual value."""
 216         value_or_array = cls._path_resolve(datablock, data_path)
 217
 218         # Both indices -1 and 0 are used for non-array properties.
 219         # -1 cannot be used in arrays, whereas 0 can be used in both arrays and non-arrays.
 220
 221         if array_index == -1:
 222             return cast(FCurveValue, value_or_array)
 223
 224         if array_index == 0:
 225             value_or_array = cls._path_resolve(datablock, data_path)
 226             try:
 227                 # MyPy doesn't understand this try/except is to determine the type.
 228                 value = value_or_array[array_index]  # type: ignore
 229             except TypeError:
 230                 # Not an array after all.
 231                 return cast(FCurveValue, value_or_array)
 232             return cast(FCurveValue, value)
 233
 234         # MyPy doesn't understand that array_index>0 implies this is indexable.
 235         return cast(FCurveValue, value_or_array[array_index])  # type: ignore
 236
 237     @staticmethod
 238     def _path_resolve(
 239         datablock: bpy.types.ID, data_path: str
 240     ) -> Union[FCurveValue, Iterable[FCurveValue]]:
 241         """Wrapper for datablock.path_resolve(data_path).
 242
 243         Raise UnresolvablePathError when the path cannot be resolved.
 244         This is easier to deal with upstream than the generic ValueError raised
 245         by Blender.
 246         """
 247         try:
 248             return datablock.path_resolve(data_path)  # type: ignore
 249         except ValueError as ex:
 250             raise UnresolvablePathError(str(ex)) from ex
 251
 252     @functools.lru_cache(maxsize=1024)
 253     def _find_bone(self, bone_name: str) -> Bone:
 254         """Find a bone by name.
 255
 256         Assumes the named bone exists, as the bones this class handles comes
 257         from the user's selection, and you can't select a non-existent bone.
 258         """
 259
 260         bone: Bone = self.params.armature_ob.pose.bones[bone_name]
 261         return bone
 262
 263     def _has_key_on_frame(self, fcurve: FCurve) -> bool:
 264         """Return True iff the FCurve has a key on the source frame."""
 265
 266         points = fcurve.keyframe_points
 267         if not points:
 268             return False
 269
 270         frame_to_find = self.params.src_frame_nr
 271         margin = 0.001
 272         high = len(points) - 1
 273         low = 0
 274         while low <= high:
 275             mid = (high + low) // 2
 276             diff = points[mid].co.x - frame_to_find
 277             if abs(diff) < margin:
 278                 return True
 279             if diff < 0:
 280                 # Frame to find is bigger than the current middle.
 281                 low = mid + 1
 282             else:
 283                 # Frame to find is smaller than the current middle
 284                 high = mid - 1
 285         return False
 286
 287
 288 def create_pose_asset(
 289     params: PoseCreationParams,
 290 ) -> Optional[Action]:
 291     """Create a single-frame Action containing only the pose of the given bones.
 292
 293     DOES mark as asset, DOES NOT configure asset metadata.
 294     """
 295
 296     creator = PoseActionCreator(params)
 297     pose_action = creator.create()
 298     if pose_action is None:
 299         return None
 300
 301     pose_action.asset_mark()
 302     pose_action.asset_generate_preview()
 303     return pose_action
 304
 305
 306 def create_pose_asset_from_context(context: Context, new_asset_name: str) -> Optional[Action]:
 307     """Create Action asset from active object & selected bones."""
 308
 309     bones = context.selected_pose_bones_from_active_object
 310     bone_names = {bone.name for bone in bones}
 311
 312     params = PoseCreationParams(
 313         context.object,
 314         getattr(context.object.animation_data, "action", None),
 315         context.scene.frame_current,
 316         frozenset(bone_names),
 317         new_asset_name,
 318     )
 319
 320     return create_pose_asset(params)
 321
 322
 323 def copy_fcurves(
 324     dst_action: Action,
 325     src_action: Action,
 326     src_frame_nr: float,
 327     bone_names: Set[str],
 328 ) -> int:
 329     """Copy FCurves, returning number of curves copied."""
 330     num_fcurves_copied = 0
 331     for fcurve in src_action.fcurves:
 332         match = pose_bone_re.match(fcurve.data_path)
 333         if not match:
 334             continue
 335
 336         bone_name = match.group(1)
 337         if bone_name not in bone_names:
 338             continue
 339
 340         # Check if there is a keyframe on this frame.
 341         keyframe = find_keyframe(fcurve, src_frame_nr)
 342         if keyframe is None:
 343             continue
 344         create_single_key_fcurve(dst_action, fcurve, keyframe)
 345         num_fcurves_copied += 1
 346     return num_fcurves_copied
 347
 348
 349 def create_single_key_fcurve(
 350     dst_action: Action, src_fcurve: FCurve, src_keyframe: Keyframe
 351 ) -> FCurve:
 352     """Create a copy of the source FCurve, but only for the given keyframe.
 353
 354     Returns a new FCurve with just one keyframe.
 355     """
 356
 357     dst_fcurve = copy_fcurve_without_keys(dst_action, src_fcurve)
 358     copy_keyframe(dst_fcurve, src_keyframe)
 359     return dst_fcurve
 360
 361
 362 def copy_fcurve_without_keys(dst_action: Action, src_fcurve: FCurve) -> FCurve:
 363     """Create a new FCurve and copy some properties."""
 364
 365     src_group_name = src_fcurve.group.name if src_fcurve.group else ""
 366     dst_fcurve = dst_action.fcurves.new(
 367         src_fcurve.data_path, index=src_fcurve.array_index, action_group=src_group_name
 368     )
 369     for propname in {"auto_smoothing", "color", "color_mode", "extrapolation"}:
 370         setattr(dst_fcurve, propname, getattr(src_fcurve, propname))
 371     return dst_fcurve
 372
 373
 374 def copy_keyframe(dst_fcurve: FCurve, src_keyframe: Keyframe) -> Keyframe:
 375     """Copy a keyframe from one FCurve to the other."""
 376
 377     dst_keyframe = dst_fcurve.keyframe_points.insert(
 378         src_keyframe.co.x, src_keyframe.co.y, options={'FAST'}, keyframe_type=src_keyframe.type
 379     )
 380
 381     for propname in {
 382         "amplitude",
 383         "back",
 384         "easing",
 385         "handle_left",
 386         "handle_left_type",
 387         "handle_right",
 388         "handle_right_type",
 389         "interpolation",
 390         "period",
 391     }:
 392         setattr(dst_keyframe, propname, getattr(src_keyframe, propname))
 393     dst_fcurve.update()
 394     return dst_keyframe
 395
 396
 397 def find_keyframe(fcurve: FCurve, frame: float) -> Optional[Keyframe]:
 398     # Binary search adapted from https://pythonguides.com/python-binary-search/
 399     keyframes = fcurve.keyframe_points
 400     low = 0
 401     high = len(keyframes) - 1
 402     mid = 0
 403
 404     # Accept any keyframe that's within 'epsilon' of the requested frame.
 405     # This should account for rounding errors and the likes.
 406     epsilon = 1e-4
 407     frame_lowerbound = frame - epsilon
 408     frame_upperbound = frame + epsilon
 409     while low <= high:
 410         mid = (high + low) // 2
 411         keyframe = keyframes[mid]
 412         if keyframe.co.x < frame_lowerbound:
 413             low = mid + 1
 414         elif keyframe.co.x > frame_upperbound:
 415             high = mid - 1
 416         else:
 417             return keyframe
 418     return None
 419
 420
 421 def assign_from_asset_browser(asset: Action, asset_browser_area: bpy.types.Area) -> None:
 422     """Assign some things from the asset browser to the asset.
 423
 424     This sets the current catalog ID, and in the future could include tags
 425     from the active dynamic catalog, etc.
 426     """
 427
 428     cat_id = asset_browser.active_catalog_id(asset_browser_area)
 429     asset.asset_data.catalog_id = cat_id