diff options
Diffstat (limited to 'rigify/base_rig.py')
| -rw-r--r-- | rigify/base_rig.py | 49 |
1 files changed, 38 insertions, 11 deletions
diff --git a/rigify/base_rig.py b/rigify/base_rig.py index b0bcc027..c284cc20 100644 --- a/rigify/base_rig.py +++ b/rigify/base_rig.py @@ -2,11 +2,11 @@ import collections -from bpy.types import PoseBone -from typing import TYPE_CHECKING, Any, Callable, Optional +from bpy.types import PoseBone, UILayout, Context +from typing import TYPE_CHECKING, Any, Callable, Optional, TypeVar, Generic from .utils.errors import RaiseErrorMixin -from .utils.bones import BoneDict, BoneUtilityMixin +from .utils.bones import BoneDict, BoneUtilityMixin, TypedBoneDict, BaseBoneDict from .utils.mechanism import MechanismUtilityMixin from .utils.metaclass import BaseStagedClass from .utils.misc import ArmatureObject @@ -144,14 +144,19 @@ class GenerateCallbackHost(BaseStagedClass, define_stages=True): pass -class BaseRig(GenerateCallbackHost, RaiseErrorMixin, BoneUtilityMixin, MechanismUtilityMixin): +_Org = TypeVar('_Org', bound=str | list[str] | BaseBoneDict) +_Ctrl = TypeVar('_Ctrl', bound=str | list[str] | BaseBoneDict) +_Mch = TypeVar('_Mch', bound=str | list[str] | BaseBoneDict) +_Deform = TypeVar('_Deform', bound=str | list[str] | BaseBoneDict) + + +class BaseRigMixin(RaiseErrorMixin, BoneUtilityMixin, MechanismUtilityMixin): generator: 'BaseGenerator' obj: ArmatureObject script: 'ScriptGenerator' base_bone: str params: Any - bones: BoneDict rigify_parent: Optional['BaseRig'] rigify_children: list['BaseRig'] @@ -160,6 +165,28 @@ class BaseRig(GenerateCallbackHost, RaiseErrorMixin, BoneUtilityMixin, Mechanism rigify_new_bones: dict[str, Optional[str]] rigify_derived_bones: dict[str, set[str]] + ############################################## + # Annotated bone containers + + class ToplevelBones(TypedBoneDict, Generic[_Org, _Ctrl, _Mch, _Deform]): + org: _Org + ctrl: _Ctrl + mch: _Mch + deform: _Deform + + class CtrlBones(TypedBoneDict): + pass + + class MchBones(TypedBoneDict): + pass + + bones: ToplevelBones[str | list[str] | BoneDict, + str | list[str] | BoneDict, # Use CtrlBones in overrides + str | list[str] | BoneDict, # Use MchBones in overrides + str | list[str] | BoneDict] + + +class BaseRig(GenerateCallbackHost, BaseRigMixin): """ Base class for all rigs. @@ -182,7 +209,7 @@ class BaseRig(GenerateCallbackHost, RaiseErrorMixin, BoneUtilityMixin, Mechanism self.params = get_rigify_params(pose_bone) # Collection of bone names for use in implementing the rig - self.bones = BoneDict( + self.bones = self.ToplevelBones( # ORG bone names org=self.find_org_bones(pose_bone), # Control bone names @@ -205,7 +232,7 @@ class BaseRig(GenerateCallbackHost, RaiseErrorMixin, BoneUtilityMixin, Mechanism self.rigify_new_bones = dict() self.rigify_derived_bones = collections.defaultdict(set) - def register_new_bone(self, new_name, old_name=None): + def register_new_bone(self, new_name: str, old_name: Optional[str] = None): """Registers this rig as the owner of this new bone.""" self.rigify_new_bones[new_name] = old_name self.generator.bone_owners[new_name] = self @@ -216,7 +243,7 @@ class BaseRig(GenerateCallbackHost, RaiseErrorMixin, BoneUtilityMixin, Mechanism ########################################################### # Bone ownership - def find_org_bones(self, pose_bone: PoseBone) -> str | list[str] | BoneDict: + def find_org_bones(self, pose_bone: PoseBone) -> str | list[str] | BaseBoneDict: """ Select bones directly owned by the rig. Returning the same bone from multiple rigs is an error. @@ -240,7 +267,7 @@ class BaseRig(GenerateCallbackHost, RaiseErrorMixin, BoneUtilityMixin, Mechanism pass @classmethod - def parameters_ui(cls, layout, params): + def parameters_ui(cls, layout: UILayout, params): """ This method draws the UI of the rigify_parameters defined on the pose_bone :param layout: @@ -250,7 +277,7 @@ class BaseRig(GenerateCallbackHost, RaiseErrorMixin, BoneUtilityMixin, Mechanism pass @classmethod - def on_parameter_update(cls, context, pose_bone, params, param_name): + def on_parameter_update(cls, context: Context, pose_bone: PoseBone, params, param_name: str): """ A callback invoked whenever a parameter value is changed by the user. """ @@ -267,7 +294,7 @@ class RigUtility(BoneUtilityMixin, MechanismUtilityMixin): self.owner = owner self.obj = owner.obj - def register_new_bone(self, new_name, old_name=None): + def register_new_bone(self, new_name: str, old_name: Optional[str] = None): self.owner.register_new_bone(new_name, old_name) |