rigify/base_rig.py

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