1 # SPDX-License-Identifier: GPL-2.0-or-later
8 from typing
import Optional
, TYPE_CHECKING
, Collection
, List
9 from bpy
.types
import PoseBone
, Bone
11 from .utils
.errors
import MetarigError
, RaiseErrorMixin
12 from .utils
.naming
import random_id
13 from .utils
.metaclass
import SingletonPluginMetaclass
14 from .utils
.rig
import list_bone_names_depth_first_sorted
, get_rigify_type
, get_rigify_params
15 from .utils
.misc
import clone_parameters
, assign_parameters
, ArmatureObject
17 from . import base_rig
19 from itertools
import count
22 from .rig_ui_template
import ScriptGenerator
25 ##############################################
27 ##############################################
30 class GeneratorPlugin(base_rig
.GenerateCallbackHost
, metaclass
=SingletonPluginMetaclass
):
32 Base class for generator plugins.
34 Generator plugins are per-Generator singleton utility
35 classes that receive the same stage callbacks as rigs.
37 Useful for building entities shared by multiple rigs
38 (e.g. the python script), or for making fire-and-forget
39 utilities that actually require multiple stages to
42 This will create only one instance per set of args:
44 instance = PluginClass(generator, ...init args)
49 def __init__(self
, generator
: 'BaseGenerator'):
50 self
.generator
= generator
51 self
.obj
= generator
.obj
53 def register_new_bone(self
, new_name
: str, old_name
: Optional
[str] = None):
54 self
.generator
.bone_owners
[new_name
] = None
56 self
.generator
.derived_bones
[old_name
].add(new_name
)
59 ##############################################
60 # Rig Substitution Mechanism
61 ##############################################
64 class SubstitutionRig(RaiseErrorMixin
):
65 """A proxy rig that replaces itself with one or more different rigs."""
67 def __init__(self
, generator
: 'BaseGenerator', pose_bone
: PoseBone
):
68 self
.generator
= generator
70 self
.obj
= generator
.obj
71 self
.base_bone
= pose_bone
.name
72 self
.params
= get_rigify_params(pose_bone
)
73 self
.params_copy
= clone_parameters(self
.params
)
76 # return [rig1, rig2...]
77 raise NotImplementedError
80 def register_new_bone(self
, new_name
: str, old_name
: Optional
[str] = None):
83 def get_params(self
, bone_name
: str):
84 return get_rigify_params(self
.obj
.pose
.bones
[bone_name
])
86 def assign_params(self
, bone_name
: str, param_dict
=None, **params
):
87 assign_parameters(self
.get_params(bone_name
), param_dict
, **params
)
89 def instantiate_rig(self
, rig_class
: str |
type, bone_name
: str):
90 if isinstance(rig_class
, str):
91 rig_class
= self
.generator
.find_rig_class(rig_class
)
93 return self
.generator
.instantiate_rig(rig_class
, self
.obj
.pose
.bones
[bone_name
])
96 ##############################################
98 ##############################################
101 class LegacyRig(base_rig
.BaseRig
):
102 """Wrapper around legacy style rigs without a common base class"""
104 def __init__(self
, generator
: 'BaseGenerator', pose_bone
: PoseBone
, wrapped_class
: type):
105 self
.wrapped_rig
= None
106 self
.wrapped_class
= wrapped_class
108 super().__init
__(generator
, pose_bone
)
110 def find_org_bones(self
, pose_bone
: PoseBone
):
111 bone_name
= pose_bone
.name
113 if not self
.wrapped_rig
:
114 self
.wrapped_rig
= self
.wrapped_class(self
.obj
, self
.base_bone
, self
.params
)
116 # Switch back to OBJECT mode if the rig changed it
117 if self
.obj
.mode
!= 'OBJECT':
118 bpy
.ops
.object.mode_set(mode
='OBJECT')
120 # Try to extract the main list of bones - old rigs often have it.
121 # This is not actually strictly necessary, so failing is OK.
122 if hasattr(self
.wrapped_rig
, 'org_bones'):
123 bones
= self
.wrapped_rig
.org_bones
124 if isinstance(bones
, list):
129 def generate_bones(self
):
130 # Inject references into the rig if it won't cause conflict
131 if not hasattr(self
.wrapped_rig
, 'rigify_generator'):
132 self
.wrapped_rig
.rigify_generator
= self
.generator
133 if not hasattr(self
.wrapped_rig
, 'rigify_wrapper'):
134 self
.wrapped_rig
.rigify_wrapper
= self
136 # Old rigs only have one generate method, so call it from
137 # generate_bones, which is the only stage allowed to add bones.
138 scripts
= self
.wrapped_rig
.generate()
140 # Switch back to EDIT mode if the rig changed it
141 if self
.obj
.mode
!= 'EDIT':
142 bpy
.ops
.object.mode_set(mode
='EDIT')
144 if isinstance(scripts
, dict):
145 if 'script' in scripts
:
146 self
.script
.add_panel_code(scripts
['script'])
147 if 'imports' in scripts
:
148 self
.script
.add_imports(scripts
['imports'])
149 if 'utilities' in scripts
:
150 self
.script
.add_utilities(scripts
['utilities'])
151 if 'register' in scripts
:
152 self
.script
.register_classes(scripts
['register'])
153 if 'register_drivers' in scripts
:
154 self
.script
.register_driver_functions(scripts
['register_drivers'])
155 if 'register_props' in scripts
:
156 for prop
, val
in scripts
['register_props']:
157 self
.script
.register_property(prop
, val
)
158 if 'noparent_bones' in scripts
:
159 for bone_name
in scripts
['noparent_bones']:
160 self
.generator
.disable_auto_parent(bone_name
)
161 elif scripts
is not None:
162 self
.script
.add_panel_code([scripts
[0]])
165 if hasattr(self
.wrapped_rig
, 'glue'):
166 self
.wrapped_rig
.glue()
168 # Switch back to OBJECT mode if the rig changed it
169 if self
.obj
.mode
!= 'OBJECT':
170 bpy
.ops
.object.mode_set(mode
='OBJECT')
173 ##############################################
174 # Base Generate Engine
175 ##############################################
179 """Base class for the main generator object. Contains rig and plugin management code."""
181 instance
: Optional
['BaseGenerator'] = None # static
183 context
: bpy
.types
.Context
184 scene
: bpy
.types
.Scene
185 view_layer
: bpy
.types
.ViewLayer
186 layer_collection
: bpy
.types
.LayerCollection
187 collection
: bpy
.types
.Collection
189 metarig
: ArmatureObject
192 script
: 'ScriptGenerator'
194 rig_list
: List
[base_rig
.BaseRig
]
195 root_rigs
: List
[base_rig
.BaseRig
]
197 bone_owners
: dict[str, Optional
[base_rig
.BaseRig
]]
198 derived_bones
: dict[str, set[str]]
203 widget_collection
: bpy
.types
.Collection
204 use_mirror_widgets
: bool
205 old_widget_table
: dict[str, bpy
.types
.Object
]
206 new_widget_table
: dict[str, bpy
.types
.Object
]
207 widget_mirror_mesh
: dict[str, bpy
.types
.Mesh
]
209 def __init__(self
, context
, metarig
):
210 self
.context
= context
211 self
.scene
= context
.scene
212 self
.view_layer
= context
.view_layer
213 self
.layer_collection
= context
.layer_collection
214 self
.collection
= self
.layer_collection
.collection
215 self
.metarig
= metarig
217 # List of all rig instances
219 # List of rigs that don't have a parent
221 # Map from bone names to their rigs
222 self
.bone_owners
= {}
223 self
.derived_bones
= collections
.defaultdict(set)
226 self
.plugin_list
= []
229 # Current execution stage so plugins could check they are used correctly
232 # Set of bones that should be left without parent
233 self
.noparent_bones
= set()
235 # Table of layer priorities for defining bone groups
236 self
.layer_group_priorities
= collections
.defaultdict(dict)
238 # Random string with time appended so that
239 # different rigs don't collide id's
240 self
.rig_id
= random_id(16)
242 # Table of renamed ORG bones
243 self
.org_rename_table
= dict()
245 def disable_auto_parent(self
, bone_name
: str):
246 """Prevent automatically parenting the bone to root if parentless."""
247 self
.noparent_bones
.add(bone_name
)
249 def find_derived_bones(self
, bone_name
: str, *, by_owner
=False, recursive
=True) -> set[str]:
250 """Find which bones were copied from the specified one."""
252 owner
= self
.bone_owners
.get(bone_name
, None)
256 table
= owner
.rigify_derived_bones
258 table
= self
.derived_bones
264 for child
in table
.get(name
, []):
272 return set(table
.get(bone_name
, []))
274 def set_layer_group_priority(self
, bone_name
: str,
275 layers
: Collection
[bool], priority
: float):
276 for i
, val
in enumerate(layers
):
278 self
.layer_group_priorities
[bone_name
][i
] = priority
280 def rename_org_bone(self
, old_name
: str, new_name
: str) -> str:
281 assert self
.stage
== 'instantiate'
282 assert old_name
== self
.org_rename_table
.get(old_name
, None)
283 assert old_name
not in self
.bone_owners
285 bone
= self
.obj
.data
.bones
[old_name
]
290 self
.org_rename_table
[old_name
] = new_name
293 def __run_object_stage(self
, method_name
: str):
294 """Run a generation stage in Object mode."""
295 assert(self
.context
.active_object
== self
.obj
)
296 assert(self
.obj
.mode
== 'OBJECT')
297 num_bones
= len(self
.obj
.data
.bones
)
299 self
.stage
= method_name
301 for rig
in self
.rig_list
:
302 rig
.rigify_invoke_stage(method_name
)
304 assert(self
.context
.active_object
== self
.obj
)
305 assert(self
.obj
.mode
== 'OBJECT')
306 assert(num_bones
== len(self
.obj
.data
.bones
))
308 # Allow plugins to be added to the end of the list on the fly
310 if i
>= len(self
.plugin_list
):
313 self
.plugin_list
[i
].rigify_invoke_stage(method_name
)
315 assert(self
.context
.active_object
== self
.obj
)
316 assert(self
.obj
.mode
== 'OBJECT')
317 assert(num_bones
== len(self
.obj
.data
.bones
))
319 def __run_edit_stage(self
, method_name
: str):
320 """Run a generation stage in Edit mode."""
321 assert(self
.context
.active_object
== self
.obj
)
322 assert(self
.obj
.mode
== 'EDIT')
323 num_bones
= len(self
.obj
.data
.edit_bones
)
325 self
.stage
= method_name
327 for rig
in self
.rig_list
:
328 rig
.rigify_invoke_stage(method_name
)
330 assert(self
.context
.active_object
== self
.obj
)
331 assert(self
.obj
.mode
== 'EDIT')
332 assert(num_bones
== len(self
.obj
.data
.edit_bones
))
334 # Allow plugins to be added to the end of the list on the fly
336 if i
>= len(self
.plugin_list
):
339 self
.plugin_list
[i
].rigify_invoke_stage(method_name
)
341 assert(self
.context
.active_object
== self
.obj
)
342 assert(self
.obj
.mode
== 'EDIT')
343 assert(num_bones
== len(self
.obj
.data
.edit_bones
))
345 def invoke_initialize(self
):
346 self
.__run
_object
_stage
('initialize')
348 def invoke_prepare_bones(self
):
349 self
.__run
_edit
_stage
('prepare_bones')
351 def __auto_register_bones(self
, bones
, rig
, plugin
=None):
352 """Find bones just added and not registered by this rig."""
355 if name
not in self
.bone_owners
:
356 self
.bone_owners
[name
] = rig
358 rig
.rigify_new_bones
[name
] = None
360 if not isinstance(rig
, LegacyRig
):
361 print(f
"WARNING: rig {self.describe_rig(rig)} "
362 f
"didn't register bone {name}\n")
364 print(f
"WARNING: plugin {plugin} didn't register bone {name}\n")
366 def invoke_generate_bones(self
):
367 assert(self
.context
.active_object
== self
.obj
)
368 assert(self
.obj
.mode
== 'EDIT')
370 self
.stage
= 'generate_bones'
372 for rig
in self
.rig_list
:
373 rig
.rigify_invoke_stage('generate_bones')
375 assert(self
.context
.active_object
== self
.obj
)
376 assert(self
.obj
.mode
== 'EDIT')
378 self
.__auto
_register
_bones
(self
.obj
.data
.edit_bones
, rig
)
380 # Allow plugins to be added to the end of the list on the fly
382 if i
>= len(self
.plugin_list
):
385 self
.plugin_list
[i
].rigify_invoke_stage('generate_bones')
387 assert(self
.context
.active_object
== self
.obj
)
388 assert(self
.obj
.mode
== 'EDIT')
390 self
.__auto
_register
_bones
(self
.obj
.data
.edit_bones
, None, plugin
=self
.plugin_list
[i
])
392 def invoke_parent_bones(self
):
393 self
.__run
_edit
_stage
('parent_bones')
395 def invoke_configure_bones(self
):
396 self
.__run
_object
_stage
('configure_bones')
398 def invoke_preapply_bones(self
):
399 self
.__run
_object
_stage
('preapply_bones')
401 def invoke_apply_bones(self
):
402 self
.__run
_edit
_stage
('apply_bones')
404 def invoke_rig_bones(self
):
405 self
.__run
_object
_stage
('rig_bones')
407 def invoke_generate_widgets(self
):
408 self
.__run
_object
_stage
('generate_widgets')
410 def invoke_finalize(self
):
411 self
.__run
_object
_stage
('finalize')
413 def instantiate_rig(self
, rig_class
: type, pose_bone
: PoseBone
) -> base_rig
.BaseRig
:
414 assert not issubclass(rig_class
, SubstitutionRig
)
416 if issubclass(rig_class
, base_rig
.BaseRig
):
417 return rig_class(self
, pose_bone
)
419 return LegacyRig(self
, pose_bone
, rig_class
)
421 def find_rig_class(self
, rig_type
: str) -> type:
422 raise NotImplementedError
424 def instantiate_rig_by_type(self
, rig_type
: str, pose_bone
: PoseBone
):
425 return self
.instantiate_rig(self
.find_rig_class(rig_type
), pose_bone
)
427 # noinspection PyMethodMayBeStatic
428 def describe_rig(self
, rig
: base_rig
.BaseRig
) -> str:
429 base_bone
= rig
.base_bone
431 if isinstance(rig
, LegacyRig
):
432 rig
= rig
.wrapped_rig
434 return "%s (%s)" % (rig
.__class
__, base_bone
)
436 def __create_rigs(self
, bone_name
, halt_on_missing
):
437 """Recursively walk bones and create rig instances."""
439 pose_bone
= self
.obj
.pose
.bones
[bone_name
]
441 rig_type
= get_rigify_type(pose_bone
)
445 rig_class
= self
.find_rig_class(rig_type
)
447 if issubclass(rig_class
, SubstitutionRig
):
448 rigs
= rig_class(self
, pose_bone
).substitute()
450 rigs
= [self
.instantiate_rig(rig_class
, pose_bone
)]
452 assert(self
.context
.active_object
== self
.obj
)
453 assert(self
.obj
.mode
== 'OBJECT')
456 self
.rig_list
.append(rig
)
458 for org_name
in rig
.rigify_org_bones
:
459 if org_name
in self
.bone_owners
:
460 old_rig
= self
.describe_rig(self
.bone_owners
[org_name
])
461 new_rig
= self
.describe_rig(rig
)
462 print(f
"CONFLICT: bone {org_name} is claimed by rigs "
463 f
"{old_rig} and {new_rig}\n")
465 self
.bone_owners
[org_name
] = rig
468 message
= f
"Rig Type Missing: python module for type '{rig_type}' "\
469 f
"not found (bone: {bone_name})"
471 raise MetarigError(message
)
474 print('print_exc():')
475 traceback
.print_exc(file=sys
.stdout
)
477 def __build_rig_tree_rec(self
, bone
: Bone
, current_rig
: Optional
[base_rig
.BaseRig
],
478 handled
: dict[base_rig
.BaseRig
, str]):
479 """Recursively walk bones and connect rig instances into a tree."""
481 rig
= self
.bone_owners
.get(bone
.name
)
484 if rig
is current_rig
:
487 elif rig
not in handled
:
488 rig
.rigify_parent
= current_rig
491 current_rig
.rigify_children
.append(rig
)
493 self
.root_rigs
.append(rig
)
495 handled
[rig
] = bone
.name
497 elif rig
.rigify_parent
is not current_rig
:
498 raise MetarigError("CONFLICT: bone {bone.name} owned by rig {rig.base_bone} "
499 f
"has different parent rig from {handled[rig]}")
504 current_rig
.rigify_child_bones
.add(bone
.name
)
506 self
.bone_owners
[bone
.name
] = current_rig
508 for child
in bone
.children
:
509 self
.__build
_rig
_tree
_rec
(child
, current_rig
, handled
)
511 def instantiate_rig_tree(self
, halt_on_missing
=False):
512 """Create rig instances and connect them into a tree."""
514 assert(self
.context
.active_object
== self
.obj
)
515 assert(self
.obj
.mode
== 'OBJECT')
517 self
.stage
= 'instantiate'
519 # Compute the list of bones
520 bone_list
= list_bone_names_depth_first_sorted(self
.obj
)
522 self
.org_rename_table
= {n
: n
for n
in bone_list
}
524 # Construct the rig instances
525 for name
in bone_list
:
526 self
.__create
_rigs
(self
.org_rename_table
[name
], halt_on_missing
)
528 # Connect rigs and bones into a tree
531 for bone
in self
.obj
.data
.bones
:
532 if bone
.parent
is None:
533 self
.__build
_rig
_tree
_rec
(bone
, None, handled
)