Update scripts to account for removal of the context override to bpy.ops
[blender-addons.git] / rigify / base_rig.py
blob682514bc9b9be10df5c0b633d118776a44917bff
1 # SPDX-License-Identifier: GPL-2.0-or-later
3 import collections
5 from bpy.types import PoseBone, UILayout, Context
6 from typing import TYPE_CHECKING, Any, Callable, Optional, TypeVar, Generic
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
15 if TYPE_CHECKING:
16 from .base_generate import BaseGenerator
17 from .rig_ui_template import ScriptGenerator
20 ##############################################
21 # Base Rig
22 ##############################################
24 class GenerateCallbackHost(BaseStagedClass, define_stages=True):
25 """
26 Standard set of callback methods to redefine.
27 Shared between BaseRig and GeneratorPlugin.
29 These callbacks are called in this order; every one is
30 called for all rigs before proceeding to the next stage.
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.
36 After each callback, all other methods decorated with
37 @stage.<method_name> are called, for instance:
39 def generate_bones(self):
40 print('first')
42 @stage.generate_bones
43 def foo(self):
44 print('second')
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.:
51 class Base(...):
52 @stage.generate_bones
53 def first(self):...
55 @stage.generate_bones
56 def second(self):...
58 class Derived(Base):
59 @stage.generate_bones
60 def third(self):...
62 # Was first defined in Base so still first:
63 @stage.generate_bones
64 def first(self):...
66 @stage.generate_bones
67 def fourth(self):...
69 Multiple inheritance can make this ordering confusing, so it
70 is best to avoid it.
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
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
90 def generate_bones(self):
91 """
92 Create all bones.
93 Called in Edit mode.
94 """
95 pass
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.
102 pass
104 def configure_bones(self):
106 Configure bone properties, e.g. transform locks, layers etc.
107 Called in Object mode. May not do Edit mode operations.
109 pass
111 def preapply_bones(self):
113 Read bone matrices for applying to edit mode.
114 Called in Object mode. May not do Edit mode operations.
116 pass
118 def apply_bones(self):
120 Can be used to apply some constraints to rest pose, and for final parenting.
121 Called in Edit mode. May not add bones.
123 pass
125 def rig_bones(self):
127 Create and configure all constraints, drivers etc.
128 Called in Object mode. May not do Edit mode operations.
130 pass
132 def generate_widgets(self):
134 Create all widget objects.
135 Called in Object mode. May not do Edit mode operations.
137 pass
139 def finalize(self):
141 Finishing touches to the construction of the rig.
142 Called in Object mode. May not do Edit mode operations.
144 pass
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)
153 class BaseRigMixin(RaiseErrorMixin, BoneUtilityMixin, MechanismUtilityMixin):
154 generator: 'BaseGenerator'
156 obj: ArmatureObject
157 script: 'ScriptGenerator'
158 base_bone: str
159 params: Any
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]]
168 ##############################################
169 # Annotated bone containers
171 class ToplevelBones(TypedBoneDict, Generic[_Org, _Ctrl, _Mch, _Deform]):
172 org: _Org
173 ctrl: _Ctrl
174 mch: _Mch
175 deform: _Deform
177 class CtrlBones(TypedBoneDict):
178 pass
180 class MchBones(TypedBoneDict):
181 pass
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]
193 class BaseRig(GenerateCallbackHost, BaseRigMixin):
195 Base class for all rigs.
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.
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.
207 def __init__(self, generator: 'BaseGenerator', pose_bone: PoseBone):
208 self.generator = generator
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)
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(),
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)
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)
247 ###########################################################
248 # Bone ownership
250 def find_org_bones(self, pose_bone: PoseBone) -> str | list[str] | BaseBoneDict:
252 Select bones directly owned by the rig. Returning the
253 same bone from multiple rigs is an error.
255 May return a single name, a list, or a BoneDict.
257 Called in Object mode, may not change the armature.
259 return [pose_bone.name]
261 ###########################################################
262 # Parameters and UI
264 @classmethod
265 def add_parameters(cls, params):
267 This method add more parameters to params
268 :param params: rigify_parameters of a pose_bone
269 :return:
271 pass
273 @classmethod
274 def parameters_ui(cls, layout: UILayout, params):
276 This method draws the UI of the rigify_parameters defined on the pose_bone
277 :param layout:
278 :param params:
279 :return:
281 pass
283 @classmethod
284 def on_parameter_update(cls, context: Context, pose_bone: PoseBone, params, param_name: str):
286 A callback invoked whenever a parameter value is changed by the user.
290 ##############################################
291 # Rig Utility
292 ##############################################
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
301 def register_new_bone(self, new_name: str, old_name: Optional[str] = None):
302 self.owner.register_new_bone(new_name, old_name)
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)
310 self.is_component_enabled = False
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)
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()
326 ##############################################
327 # Rig Stage Decorators
328 ##############################################
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