asset_library/pose/pose_creation.py

424 lines
15 KiB
Python

# SPDX-License-Identifier: GPL-2.0-or-later
"""
Pose Library - creation functions.
"""
import dataclasses
import functools
import re
from typing import Optional, FrozenSet, Set, Union, Iterable, cast
import bpy
from bpy.types import (
Action,
Bone,
Context,
FCurve,
Keyframe,
)
FCurveValue = Union[float, int]
pose_bone_re = re.compile(r'pose.bones\["([^"]+)"\]')
"""RegExp for matching FCurve data paths."""
@dataclasses.dataclass(unsafe_hash=True, frozen=True)
class PoseCreationParams:
armature_ob: bpy.types.Object
src_action: Optional[Action]
src_frame_nr: float
bone_names: FrozenSet[str]
new_asset_name: str
class UnresolvablePathError(ValueError):
"""Raised when a data_path cannot be resolved to a current value."""
@dataclasses.dataclass(unsafe_hash=True)
class PoseActionCreator:
"""Create an Action that's suitable for marking as Asset.
Does not mark as asset yet, nor does it add asset metadata.
"""
params: PoseCreationParams
# These were taken from Blender's Action baking code in `anim_utils.py`.
# Items are (name, array_length) tuples.
_bbone_props = [
("bbone_curveinx", None),
("bbone_curveoutx", None),
("bbone_curveinz", None),
("bbone_curveoutz", None),
("bbone_rollin", None),
("bbone_rollout", None),
("bbone_scalein", 3),
("bbone_scaleout", 3),
("bbone_easein", None),
("bbone_easeout", None),
]
def create(self) -> Optional[Action]:
"""Create a single-frame Action containing only the given bones, or None if no anim data was found."""
try:
dst_action = self._create_new_action()
self._store_pose(dst_action)
finally:
# Prevent next instantiations of this class from reusing pointers to
# bones. They may not be valid by then any more.
self._find_bone.cache_clear()
if len(dst_action.fcurves) == 0:
bpy.data.actions.remove(dst_action)
return None
return dst_action
def _create_new_action(self) -> Action:
dst_action = bpy.data.actions.new(self.params.new_asset_name)
if self.params.src_action:
dst_action.id_root = self.params.src_action.id_root
dst_action.user_clear() # actions.new() sets users=1, but marking as asset also increments user count.
return dst_action
def _store_pose(self, dst_action: Action) -> None:
"""Store the current pose into the given action."""
self._store_bone_pose_parameters(dst_action)
self._store_animated_parameters(dst_action)
self._store_parameters_from_callback(dst_action)
def _store_bone_pose_parameters(self, dst_action: Action) -> None:
"""Store loc/rot/scale/bbone values in the Action."""
for bone_name in sorted(self.params.bone_names):
self._store_location(dst_action, bone_name)
self._store_rotation(dst_action, bone_name)
self._store_scale(dst_action, bone_name)
self._store_bbone(dst_action, bone_name)
def _store_animated_parameters(self, dst_action: Action) -> None:
"""Store the current value of any animated bone properties."""
if self.params.src_action is None:
return
armature_ob = self.params.armature_ob
for fcurve in self.params.src_action.fcurves:
match = pose_bone_re.match(fcurve.data_path)
if not match:
# Not animating a bone property.
continue
bone_name = match.group(1)
if bone_name not in self.params.bone_names:
# Bone is not our export set.
continue
if dst_action.fcurves.find(fcurve.data_path, index=fcurve.array_index):
# This property is already handled by a previous _store_xxx() call.
continue
# Only include in the pose if there is a key on this frame.
if not self._has_key_on_frame(fcurve):
continue
try:
value = self._current_value(armature_ob, fcurve.data_path, fcurve.array_index)
except UnresolvablePathError:
# A once-animated property no longer exists.
continue
dst_fcurve = dst_action.fcurves.new(
fcurve.data_path, index=fcurve.array_index, action_group=bone_name
)
dst_fcurve.keyframe_points.insert(self.params.src_frame_nr, value=value)
dst_fcurve.update()
def _store_parameters_from_callback(self, dst_action: Action) -> None:
"""Store extra parameters in the pose based on arbitrary callbacks.
Not implemented yet, needs a proper design & some user stories.
"""
pass
def _store_location(self, dst_action: Action, bone_name: str) -> None:
"""Store bone location."""
self._store_bone_array(dst_action, bone_name, "location", 3)
def _store_rotation(self, dst_action: Action, bone_name: str) -> None:
"""Store bone rotation given current rotation mode."""
bone = self._find_bone(bone_name)
if bone.rotation_mode == "QUATERNION":
self._store_bone_array(dst_action, bone_name, "rotation_quaternion", 4)
elif bone.rotation_mode == "AXIS_ANGLE":
self._store_bone_array(dst_action, bone_name, "rotation_axis_angle", 4)
else:
self._store_bone_array(dst_action, bone_name, "rotation_euler", 3)
def _store_scale(self, dst_action: Action, bone_name: str) -> None:
"""Store bone scale."""
self._store_bone_array(dst_action, bone_name, "scale", 3)
def _store_bbone(self, dst_action: Action, bone_name: str) -> None:
"""Store bendy-bone parameters."""
for prop_name, array_length in self._bbone_props:
if array_length:
self._store_bone_array(dst_action, bone_name, prop_name, array_length)
else:
self._store_bone_property(dst_action, bone_name, prop_name)
def _store_bone_array(
self, dst_action: Action, bone_name: str, property_name: str, array_length: int
) -> None:
"""Store all elements of an array property."""
for array_index in range(array_length):
self._store_bone_property(dst_action, bone_name, property_name, array_index)
def _store_bone_property(
self,
dst_action: Action,
bone_name: str,
property_path: str,
array_index: int = -1,
) -> None:
"""Store the current value of a single bone property."""
bone = self._find_bone(bone_name)
value = self._current_value(bone, property_path, array_index)
# Get the full 'pose.bones["bone_name"].blablabla' path suitable for FCurves.
rna_path = bone.path_from_id(property_path)
fcurve: Optional[FCurve] = dst_action.fcurves.find(rna_path, index=array_index)
if fcurve is None:
fcurve = dst_action.fcurves.new(rna_path, index=array_index, action_group=bone_name)
fcurve.keyframe_points.insert(self.params.src_frame_nr, value=value)
fcurve.update()
@classmethod
def _current_value(
cls, datablock: bpy.types.ID, data_path: str, array_index: int
) -> FCurveValue:
"""Resolve an RNA path + array index to an actual value."""
value_or_array = cls._path_resolve(datablock, data_path)
# Both indices -1 and 0 are used for non-array properties.
# -1 cannot be used in arrays, whereas 0 can be used in both arrays and non-arrays.
if array_index == -1:
return cast(FCurveValue, value_or_array)
if array_index == 0:
value_or_array = cls._path_resolve(datablock, data_path)
try:
# MyPy doesn't understand this try/except is to determine the type.
value = value_or_array[array_index] # type: ignore
except TypeError:
# Not an array after all.
return cast(FCurveValue, value_or_array)
return cast(FCurveValue, value)
# MyPy doesn't understand that array_index>0 implies this is indexable.
return cast(FCurveValue, value_or_array[array_index]) # type: ignore
@staticmethod
def _path_resolve(
datablock: bpy.types.ID, data_path: str
) -> Union[FCurveValue, Iterable[FCurveValue]]:
"""Wrapper for datablock.path_resolve(data_path).
Raise UnresolvablePathError when the path cannot be resolved.
This is easier to deal with upstream than the generic ValueError raised
by Blender.
"""
try:
return datablock.path_resolve(data_path) # type: ignore
except ValueError as ex:
raise UnresolvablePathError(str(ex)) from ex
@functools.lru_cache(maxsize=1024)
def _find_bone(self, bone_name: str) -> Bone:
"""Find a bone by name.
Assumes the named bone exists, as the bones this class handles comes
from the user's selection, and you can't select a non-existent bone.
"""
bone: Bone = self.params.armature_ob.pose.bones[bone_name]
return bone
def _has_key_on_frame(self, fcurve: FCurve) -> bool:
"""Return True iff the FCurve has a key on the source frame."""
points = fcurve.keyframe_points
if not points:
return False
frame_to_find = self.params.src_frame_nr
margin = 0.001
high = len(points) - 1
low = 0
while low <= high:
mid = (high + low) // 2
diff = points[mid].co.x - frame_to_find
if abs(diff) < margin:
return True
if diff < 0:
# Frame to find is bigger than the current middle.
low = mid + 1
else:
# Frame to find is smaller than the current middle
high = mid - 1
return False
def create_pose_asset(
params: PoseCreationParams,
) -> Optional[Action]:
"""Create a single-frame Action containing only the pose of the given bones.
DOES mark as asset, DOES NOT configure asset metadata.
"""
creator = PoseActionCreator(params)
pose_action = creator.create()
if pose_action is None:
return None
pose_action.asset_mark()
pose_action.asset_generate_preview()
return pose_action
#def create_pose_asset_from_context(context: Context, new_asset_name: str, selection=True) -> Optional[Action]:
def create_pose_asset_from_context(context: Context, new_asset_name: str) -> Optional[Action]:
"""Create Action asset from active object & selected bones."""
bones = context.selected_pose_bones_from_active_object
bone_names = {bone.name for bone in bones}
params = PoseCreationParams(
context.object,
getattr(context.object.animation_data, "action", None),
context.scene.frame_current,
frozenset(bone_names),
new_asset_name,
)
return create_pose_asset(params)
def copy_fcurves(
dst_action: Action,
src_action: Action,
src_frame_nr: float,
bone_names: Set[str],
) -> int:
"""Copy FCurves, returning number of curves copied."""
num_fcurves_copied = 0
for fcurve in src_action.fcurves:
match = pose_bone_re.match(fcurve.data_path)
if not match:
continue
bone_name = match.group(1)
if bone_name not in bone_names:
continue
# Check if there is a keyframe on this frame.
keyframe = find_keyframe(fcurve, src_frame_nr)
if keyframe is None:
continue
create_single_key_fcurve(dst_action, fcurve, keyframe)
num_fcurves_copied += 1
return num_fcurves_copied
def create_single_key_fcurve(
dst_action: Action, src_fcurve: FCurve, src_keyframe: Keyframe
) -> FCurve:
"""Create a copy of the source FCurve, but only for the given keyframe.
Returns a new FCurve with just one keyframe.
"""
dst_fcurve = copy_fcurve_without_keys(dst_action, src_fcurve)
copy_keyframe(dst_fcurve, src_keyframe)
return dst_fcurve
def copy_fcurve_without_keys(dst_action: Action, src_fcurve: FCurve) -> FCurve:
"""Create a new FCurve and copy some properties."""
src_group_name = src_fcurve.group.name if src_fcurve.group else ""
dst_fcurve = dst_action.fcurves.new(
src_fcurve.data_path, index=src_fcurve.array_index, action_group=src_group_name
)
for propname in {"auto_smoothing", "color", "color_mode", "extrapolation"}:
setattr(dst_fcurve, propname, getattr(src_fcurve, propname))
return dst_fcurve
def copy_keyframe(dst_fcurve: FCurve, src_keyframe: Keyframe) -> Keyframe:
"""Copy a keyframe from one FCurve to the other."""
dst_keyframe = dst_fcurve.keyframe_points.insert(
src_keyframe.co.x, src_keyframe.co.y, options={'FAST'}, keyframe_type=src_keyframe.type
)
for propname in {
"amplitude",
"back",
"easing",
"handle_left",
"handle_left_type",
"handle_right",
"handle_right_type",
"interpolation",
"period",
}:
setattr(dst_keyframe, propname, getattr(src_keyframe, propname))
dst_fcurve.update()
return dst_keyframe
def find_keyframe(fcurve: FCurve, frame: float) -> Optional[Keyframe]:
# Binary search adapted from https://pythonguides.com/python-binary-search/
keyframes = fcurve.keyframe_points
low = 0
high = len(keyframes) - 1
mid = 0
# Accept any keyframe that's within 'epsilon' of the requested frame.
# This should account for rounding errors and the likes.
epsilon = 1e-4
frame_lowerbound = frame - epsilon
frame_upperbound = frame + epsilon
while low <= high:
mid = (high + low) // 2
keyframe = keyframes[mid]
if keyframe.co.x < frame_lowerbound:
low = mid + 1
elif keyframe.co.x > frame_upperbound:
high = mid - 1
else:
return keyframe
return None
def assign_from_asset_browser(asset: Action, asset_browser_area: bpy.types.Area) -> None:
"""Assign some things from the asset browser to the asset.
This sets the current catalog ID, and in the future could include tags
from the active dynamic catalog, etc.
"""
cat_id = asset_browser.active_catalog_id(asset_browser_area)
asset.asset_data.catalog_id = cat_id