rigify/base_rig.py

   1 # SPDX-License-Identifier: GPL-2.0-or-later
   2
   3 import collections
   4
   5 from bpy.types import PoseBone, UILayout, Context
   6 from typing import TYPE_CHECKING, Any, Callable, Optional, TypeVar, Generic
   7
   8 from .utils.errors import RaiseErrorMixin
   9 from .utils.bones import BoneDict, BoneUtilityMixin, TypedBoneDict, BaseBoneDict
  10 from .utils.mechanism import MechanismUtilityMixin
  11 from .utils.metaclass import BaseStagedClass
  12 from .utils.misc import ArmatureObject
  13 from .utils.rig import get_rigify_params
  14
  15 if TYPE_CHECKING:
  16     from .base_generate import BaseGenerator
  17     from .rig_ui_template import ScriptGenerator
  18
  19
  20 ##############################################
  21 # Base Rig
  22 ##############################################
  23
  24 class GenerateCallbackHost(BaseStagedClass, define_stages=True):
  25     """
  26     Standard set of callback methods to redefine.
  27     Shared between BaseRig and GeneratorPlugin.
  28
  29     These callbacks are called in this order; every one is
  30     called for all rigs before proceeding to the next stage.
  31
  32     Switching modes is not allowed in rigs for performance
  33     reasons. Place code in the appropriate callbacks to use
  34     the mode set by the main engine.
  35
  36     After each callback, all other methods decorated with
  37     @stage.<method_name> are called, for instance:
  38
  39     def generate_bones(self):
  40         print('first')
  41
  42     @stage.generate_bones
  43     def foo(self):
  44         print('second')
  45
  46     Will print 'first', then 'second'. Multiple methods in the
  47     same stage are called in the order they are first defined;
  48     in case of inheritance, the class bodies are scanned in
  49     reverse MRO order. E.g.:
  50
  51     class Base(...):
  52         @stage.generate_bones
  53         def first(self):...
  54
  55         @stage.generate_bones
  56         def second(self):...
  57
  58     class Derived(Base):
  59         @stage.generate_bones
  60         def third(self):...
  61
  62         # Was first defined in Base so still first:
  63         @stage.generate_bones
  64         def first(self):...
  65
  66         @stage.generate_bones
  67         def fourth(self):...
  68
  69     Multiple inheritance can make this ordering confusing, so it
  70     is best to avoid it.
  71
  72     When overriding such methods in a subclass the appropriate
  73     decorator should be repeated for code clarity reasons;
  74     a warning is printed if this is not done.
  75     """
  76     def initialize(self):
  77         """
  78         Initialize processing after all rig classes are constructed.
  79         Called in Object mode. May not change the armature.
  80         """
  81         pass
  82
  83     def prepare_bones(self):
  84         """
  85         Prepare ORG bones for generation, e.g. align them.
  86         Called in Edit mode. May not add bones.
  87         """
  88         pass
  89
  90     def generate_bones(self):
  91         """
  92         Create all bones.
  93         Called in Edit mode.
  94         """
  95         pass
  96
  97     def parent_bones(self):
  98         """
  99         Parent all bones and set other edit mode properties.
 100         Called in Edit mode. May not add bones.
 101         """
 102         pass
 103
 104     def configure_bones(self):
 105         """
 106         Configure bone properties, e.g. transform locks, layers etc.
 107         Called in Object mode. May not do Edit mode operations.
 108         """
 109         pass
 110
 111     def preapply_bones(self):
 112         """
 113         Read bone matrices for applying to edit mode.
 114         Called in Object mode. May not do Edit mode operations.
 115         """
 116         pass
 117
 118     def apply_bones(self):
 119         """
 120         Can be used to apply some constraints to rest pose, and for final parenting.
 121         Called in Edit mode. May not add bones.
 122         """
 123         pass
 124
 125     def rig_bones(self):
 126         """
 127         Create and configure all constraints, drivers etc.
 128         Called in Object mode. May not do Edit mode operations.
 129         """
 130         pass
 131
 132     def generate_widgets(self):
 133         """
 134         Create all widget objects.
 135         Called in Object mode. May not do Edit mode operations.
 136         """
 137         pass
 138
 139     def finalize(self):
 140         """
 141         Finishing touches to the construction of the rig.
 142         Called in Object mode. May not do Edit mode operations.
 143         """
 144         pass
 145
 146
 147 _Org = TypeVar('_Org', bound=str | list[str] | BaseBoneDict)
 148 _Ctrl = TypeVar('_Ctrl', bound=str | list[str] | BaseBoneDict)
 149 _Mch = TypeVar('_Mch', bound=str | list[str] | BaseBoneDict)
 150 _Deform = TypeVar('_Deform', bound=str | list[str] | BaseBoneDict)
 151
 152
 153 class BaseRigMixin(RaiseErrorMixin, BoneUtilityMixin, MechanismUtilityMixin):
 154     generator: 'BaseGenerator'
 155
 156     obj: ArmatureObject
 157     script: 'ScriptGenerator'
 158     base_bone: str
 159     params: Any
 160
 161     rigify_parent: Optional['BaseRig']
 162     rigify_children: list['BaseRig']
 163     rigify_org_bones: set[str]
 164     rigify_child_bones: set[str]
 165     rigify_new_bones: dict[str, Optional[str]]
 166     rigify_derived_bones: dict[str, set[str]]
 167
 168     ##############################################
 169     # Annotated bone containers
 170
 171     class ToplevelBones(TypedBoneDict, Generic[_Org, _Ctrl, _Mch, _Deform]):
 172         org: _Org
 173         ctrl: _Ctrl
 174         mch: _Mch
 175         deform: _Deform
 176
 177     class CtrlBones(TypedBoneDict):
 178         pass
 179
 180     class MchBones(TypedBoneDict):
 181         pass
 182
 183     # Subclass and use the above CtrlBones and MchBones classes in overrides.
 184     # It is necessary to reference them via absolute strings, e.g. 'Rig.CtrlBones',
 185     # because when using just CtrlBones the annotation won't work fully in subclasses
 186     # of the rig class in PyCharm (no warnings about unknown attribute access).
 187     bones: ToplevelBones[str | list[str] | BoneDict,
 188                          str | list[str] | BoneDict,
 189                          str | list[str] | BoneDict,
 190                          str | list[str] | BoneDict]
 191
 192
 193 class BaseRig(GenerateCallbackHost, BaseRigMixin):
 194     """
 195     Base class for all rigs.
 196
 197     The main weak areas in the legacy (pre-2.76b) Rigify rig class structure
 198     was that there were no provisions for intelligent interactions
 199     between rigs, and all processing was done via one generate
 200     method, necessitating frequent expensive mode switches.
 201
 202     This structure fixes those problems by providing a mandatory
 203     base class that hold documented connections between rigs, bones,
 204     and the common generator object. The generation process is also
 205     split into multiple stages.
 206     """
 207     def __init__(self, generator: 'BaseGenerator', pose_bone: PoseBone):
 208         self.generator = generator
 209
 210         self.obj = generator.obj
 211         self.script = generator.script
 212         self.base_bone = pose_bone.name
 213         self.params = get_rigify_params(pose_bone)
 214
 215         # Collection of bone names for use in implementing the rig
 216         self.bones = self.ToplevelBones(
 217             # ORG bone names
 218             org=self.find_org_bones(pose_bone),
 219             # Control bone names
 220             ctrl=BoneDict(),
 221             # MCH bone names
 222             mch=BoneDict(),
 223             # DEF bone names
 224             deform=BoneDict(),
 225         )
 226
 227         # Data useful for complex rig interaction:
 228         # Parent-child links between rigs.
 229         self.rigify_parent = None
 230         self.rigify_children = []
 231         # ORG bones directly owned by the rig.
 232         self.rigify_org_bones = set(self.bones.flatten('org'))
 233         # Children of bones owned by the rig.
 234         self.rigify_child_bones = set()
 235         # Bones created by the rig (mapped to original names)
 236         self.rigify_new_bones = dict()
 237         self.rigify_derived_bones = collections.defaultdict(set)
 238
 239     def register_new_bone(self, new_name: str, old_name: Optional[str] = None):
 240         """Registers this rig as the owner of this new bone."""
 241         self.rigify_new_bones[new_name] = old_name
 242         self.generator.bone_owners[new_name] = self
 243         if old_name:
 244             self.rigify_derived_bones[old_name].add(new_name)
 245             self.generator.derived_bones[old_name].add(new_name)
 246
 247     ###########################################################
 248     # Bone ownership
 249
 250     def find_org_bones(self, pose_bone: PoseBone) -> str | list[str] | BaseBoneDict:
 251         """
 252         Select bones directly owned by the rig. Returning the
 253         same bone from multiple rigs is an error.
 254
 255         May return a single name, a list, or a BoneDict.
 256
 257         Called in Object mode, may not change the armature.
 258         """
 259         return [pose_bone.name]
 260
 261     ###########################################################
 262     # Parameters and UI
 263
 264     @classmethod
 265     def add_parameters(cls, params):
 266         """
 267         This method add more parameters to params
 268         :param params: rigify_parameters of a pose_bone
 269         :return:
 270         """
 271         pass
 272
 273     @classmethod
 274     def parameters_ui(cls, layout: UILayout, params):
 275         """
 276         This method draws the UI of the rigify_parameters defined on the pose_bone
 277         :param layout:
 278         :param params:
 279         :return:
 280         """
 281         pass
 282
 283     @classmethod
 284     def on_parameter_update(cls, context: Context, pose_bone: PoseBone, params, param_name: str):
 285         """
 286         A callback invoked whenever a parameter value is changed by the user.
 287         """
 288
 289
 290 ##############################################
 291 # Rig Utility
 292 ##############################################
 293
 294
 295 class RigUtility(BoneUtilityMixin, MechanismUtilityMixin):
 296     """Base class for utility classes that generate part of a rig."""
 297     def __init__(self, owner):
 298         self.owner = owner
 299         self.obj = owner.obj
 300
 301     def register_new_bone(self, new_name: str, old_name: Optional[str] = None):
 302         self.owner.register_new_bone(new_name, old_name)
 303
 304
 305 class LazyRigComponent(GenerateCallbackHost, RigUtility):
 306     """Base class for utility classes that generate part of a rig using callbacks. Starts as disabled."""
 307     def __init__(self, owner):
 308         super().__init__(owner)
 309
 310         self.is_component_enabled = False
 311
 312     def enable_component(self):
 313         if not self.is_component_enabled:
 314             self.is_component_enabled = True
 315             self.owner.rigify_sub_objects = objects = self.owner.rigify_sub_objects or []
 316             objects.append(self)
 317
 318
 319 class RigComponent(LazyRigComponent):
 320     """Base class for utility classes that generate part of a rig using callbacks."""
 321     def __init__(self, owner):
 322         super().__init__(owner)
 323         self.enable_component()
 324
 325
 326 ##############################################
 327 # Rig Stage Decorators
 328 ##############################################
 329
 330 # noinspection PyPep8Naming
 331 @GenerateCallbackHost.stage_decorator_container
 332 class stage:
 333     """Contains @stage.<...> decorators for all valid stages."""
 334     # Declare stages for auto-completion - doesn't affect execution.
 335     initialize: Callable
 336     prepare_bones: Callable
 337     generate_bones: Callable
 338     parent_bones: Callable
 339     configure_bones: Callable
 340     preapply_bones: Callable
 341     apply_bones: Callable
 342     rig_bones: Callable
 343     generate_widgets: Callable
 344     finalize: Callable