"""Armature object for managing pose"""
import numpy as np
from .bsyn_object import BsynObject, animatable_property
from .other_objects import Empty
from .utils import SelectObjects, SetMode
from typing import Union
from ..utils import types
import bpy
import mathutils
[docs]
class PoseBone(BsynObject):
def __init__(self, object: bpy.types.PoseBone):
self._object = object
object.rotation_mode = "XYZ" # all rotations done in XYZ Euler mode
@property
def name(self) -> str:
return self.obj.name
@property
def constraints(self):
return self.obj.constraints
@property
def armature(self):
return self.obj.id_data
@property
def matrix_world(self):
return self.armature.matrix_world @ self.obj.matrix
@property
def location(self) -> mathutils.Vector:
"""
Gets the world location of the head (start) of the bone.
"""
return self.obj.head
@property
def head_location(self) -> mathutils.Vector:
"""
Gets the world location of the head (start) of the bone.
"""
return self.armature.matrix_world @ self.obj.head
@property
def tail_location(self) -> mathutils.Vector:
"""
Gets the world location of the tail (or end) of the bone.
"""
return self.armature.matrix_world @ self.obj.tail
@property
def rotation_euler(self) -> mathutils.Euler:
"""Rotation of pose bone in world space"""
return self.matrix_world.to_euler()
@property
def local_rotation_euler(self) -> mathutils.Euler:
"""Rotation of pose bone in local space (ignoring parent transforms)"""
return self.obj.rotation_euler
@local_rotation_euler.setter
def local_rotation_euler(self, value: mathutils.Euler):
self.obj.rotation_euler = value
[docs]
def apply_global_rotation(self, axis: str = "X", val: float = 0.0, degrees=False):
"""Equivalent to selecting the bone in Blender, and applying a global rotation about an axis
:param axis: X, Y, or Z
:param val: Rotation amount in radians
:param degrees: If true, rotation is given in degrees"""
if degrees:
val = np.deg2rad(val)
with SetMode("POSE", object=self.armature):
with SelectObjects([self.obj]):
bpy.context.object.data.bones.active = self.armature.data.bones[
self.name
]
bpy.ops.transform.rotate(
value=val,
orient_axis=axis,
orient_type="GLOBAL",
constraint_axis=[axis == "X", axis == "Y", axis == "Z"],
)
@animatable_property("scale")
def set_scale(self, scale: types.VectorLikeOrScalar, update: bool = True):
"""Set scale of pose bone.
:param scale: Scale to set. Either single value or 3 long vector
:param update: Update the scene after setting
scale. For batch operations, can be faster to set this to False, and call `bsyn.context.view_layer.update()`
after all operations.
"""
if isinstance(scale, (int, float)):
scale = (scale, scale, scale)
self.obj.scale = scale
if update:
bpy.context.view_layer.update()
[docs]
class BoneConstraint(BsynObject):
"""Generic pose bone constraint object"""
def __init__(
self, pose_bone: PoseBone, constraint: bpy.types.Constraint, empty: Empty
):
self.pose_bone = pose_bone
self.constraint = constraint
self.empty = empty
self._object = empty.obj
@property
def name(self):
return f"{self.constraint.name}: {self.pose_bone.name} -> {self.empty.name}"
@property
def bone(self):
return self.pose_bone
[docs]
def remove(self):
"""Remove this constraint."""
self.bone.constraints.remove(self.constraint)
if self.empty is not None:
self.empty.remove()
@property
def armature(self):
return self.bone.id_data
[docs]
class Armature(BsynObject):
"""This class manages armatures - a collection of posable bones."""
def __init__(self, object: bpy.types.Armature):
self._object = object
self.ik_constraints = {} # for IK constraints
self.constraints = {} # for generic constraints
self.pose_bones = {n: PoseBone(b) for n, b in object.pose.bones.items()}
@property
def name(self) -> str:
return self.obj.name
@property
def pose(self) -> bpy.types.Pose:
return self.obj.pose
[docs]
def get_bone(self, bone: Union[str, PoseBone, bpy.types.PoseBone]) -> PoseBone:
"""
Get bone from armature.
:param bone_name: Name of bone to get (or PoseBone object)
:return: PoseBone object
"""
if isinstance(bone, PoseBone):
return bone
if isinstance(bone, bpy.types.PoseBone):
return self.pose_bones[bone.name]
if isinstance(bone, str):
return self.pose_bones[bone]
raise TypeError(
f"Expected bone to be PoseBone, str, or PoseBone, got {type(bone)}"
)
[docs]
def pose_bone(
self,
bone: Union[str, bpy.types.PoseBone],
rotation: types.VectorLike = None,
location: types.VectorLike = None,
scale: types.VectorLike = None,
frame: int = None,
):
"""Set the pose of a bone by giving a Euler XYZ rotation and/or location.
:param bone: Name of bone to pose, or PoseBone object
:param rotation: Euler XYZ rotation in radians
:param location: Location of bone
:param scale: Scale of bone
:param frame: Frame to set pose on. If given, will insert keyframe here.
"""
with SetMode("POSE", object=self.obj):
if isinstance(bone, str):
bone = self.get_bone(bone)
bone.rotation_mode = "XYZ"
if rotation is not None:
bone.set_rotation_euler(rotation, frame=frame)
if location is not None:
bone.set_location(location, frame=frame)
if scale is not None:
bone.set_scale(scale, frame=frame)
[docs]
def clear_pose(self, rot=True, location=True, scale=True, bones=None):
"""Clear the pose of the armature.
For the target bones, sets poses to zero, and removes any IK constraints.
:param rot: Clear rotation
:param location: Clear location
:param scale: Clear scale
:param bones: List of bones to clear. If not given, will clear all bones.
"""
with SetMode("POSE", object=self.obj):
for bone in self.pose.bones:
if bones is None or bone.name in bones:
if rot:
bone.rotation_euler = (0, 0, 0)
if location:
bone.location = (0, 0, 0)
if scale:
bone.scale = (1, 1, 1)
for bone_name in self.ik_constraints.keys():
if bone_name == bone.name:
self.ik_constraints[bone_name].remove() # remove the constraint
self.ik_constraints.pop(bone_name) # remove reference to it
for cname in [*self.constraints.keys()]:
constraint = self.constraints[cname]
if constraint.pose_bone.name == bone.name:
self.constraints[cname].remove()
self.constraints.pop(cname)
[docs]
def add_constraint(
self,
bone: Union[str, PoseBone],
constraint_name: str,
target: Empty = None,
**kwargs,
) -> BoneConstraint:
"""
Applies a generic bone constraint. Returns BoneConstraint object
:param bone: Bone to apply constraint to
:param constraint_name:
:param target: Empty target. If none given, will create one.
Any other constraint specific keyword arguments given as well will be fed into the new constraint
"""
bone = self.get_bone(bone)
# Add IK constraint to the specified bone
constraint = bone.constraints.new(constraint_name)
for k, v in kwargs.items():
setattr(constraint, k, v)
# Create an empty to serve as the IK target
if target is None:
target = Empty(name="Target_For_" + bone.name)
constraint.target = target.object
bpy.context.view_layer.update()
constraint = BoneConstraint(bone, constraint, target)
self.constraints[constraint.name] = constraint
return constraint