1import bpy 2 3import re 4import itertools 5import bisect 6import json 7 8from .errors import MetarigError 9from .naming import strip_prefix, make_derived_name 10from .mechanism import MechanismUtilityMixin 11from .misc import map_list, map_apply, force_lazy 12 13from ..base_rig import * 14from ..base_generate import GeneratorPlugin 15 16from collections import defaultdict 17from itertools import count, repeat, chain 18 19 20def _rig_is_child(rig, parent): 21 if parent is None: 22 return True 23 24 while rig: 25 if rig is parent: 26 return True 27 28 rig = rig.rigify_parent 29 30 return False 31 32 33class SwitchParentBuilder(GeneratorPlugin, MechanismUtilityMixin): 34 """ 35 Implements centralized generation of switchable parent mechanisms. 36 Allows all rigs to register their bones as possible parents for other rigs. 37 """ 38 39 def __init__(self, generator): 40 super().__init__(generator) 41 42 self.child_list = [] 43 self.global_parents = [] 44 self.local_parents = defaultdict(list) 45 self.child_map = {} 46 self.frozen = False 47 48 self.register_parent(None, 'root', name='Root', is_global=True) 49 50 51 ############################## 52 # API 53 54 def register_parent(self, rig, bone, *, name=None, is_global=False, exclude_self=False, inject_into=None, tags=None): 55 """ 56 Registers a bone of the specified rig as a possible parent. 57 58 Parameters: 59 rig Owner of the bone. 60 bone Actual name of the parent bone. 61 name Name of the parent for mouse-over hint. 62 is_global The parent is accessible to all rigs, instead of just children of owner. 63 exclude_self The parent is invisible to the owner rig itself. 64 inject_into Make this parent available to children of the specified rig. 65 tags Set of tags to use for default parent selection. 66 67 Lazy creation: 68 The bone parameter may be a function creating the bone on demand and 69 returning its name. It is guaranteed to be called at most once. 70 """ 71 72 assert not self.frozen 73 assert isinstance(bone, str) or callable(bone) 74 assert callable(bone) or _rig_is_child(rig, self.generator.bone_owners[bone]) 75 assert _rig_is_child(rig, inject_into) 76 77 real_rig = rig 78 79 if inject_into and inject_into is not rig: 80 rig = inject_into 81 tags = (tags or set()) | {'injected'} 82 83 entry = { 84 'rig': rig, 'bone': bone, 'name': name, 'tags': tags, 85 'is_global': is_global, 'exclude_self': exclude_self, 86 'real_rig': real_rig, 'used': False, 87 } 88 89 if is_global: 90 self.global_parents.append(entry) 91 else: 92 self.local_parents[id(rig)].append(entry) 93 94 95 def build_child(self, rig, bone, *, use_parent_mch=True, **options): 96 """ 97 Build a switchable parent mechanism for the specified bone. 98 99 Parameters: 100 rig Owner of the child bone. 101 bone Name of the child bone. 102 extra_parents List of bone names or (name, user_name) pairs to use as additional parents. 103 use_parent_mch Create an intermediate MCH bone for the constraints and parent the child to it. 104 select_parent Select the specified bone instead of the last one. 105 select_tags List of parent tags to try for default selection. 106 ignore_global Ignore the is_global flag of potential parents. 107 exclude_self Ignore parents registered by the rig itself. 108 allow_self Ignore the 'exclude_self' setting of the parent. 109 context_rig Rig to use for selecting parents; defaults to rig. 110 no_implicit Only use parents listed as extra_parents. 111 only_selected Like no_implicit, but allow the 'default' selected parent. 112 113 prop_bone Name of the bone to add the property to. 114 prop_id Actual name of the control property. 115 prop_name Name of the property to use in the UI script. 116 controls Collection of controls to bind property UI to. 117 118 ctrl_bone User visible control bone that depends on this parent (for switch & keep transform) 119 no_fix_* Disable "Switch and Keep Transform" correction for specific channels. 120 copy_* Override the specified components by copying from another bone. 121 inherit_scale Inherit scale mode for the child bone (default: AVERAGE). 122 123 Lazy parameters: 124 'extra_parents', 'select_parent', 'prop_bone', 'controls', 'copy_*' 125 may be a function returning the value. They are called in the configure_bones stage. 126 """ 127 assert self.generator.stage == 'generate_bones' and not self.frozen 128 assert rig is not None 129 assert isinstance(bone, str) 130 assert bone not in self.child_map 131 132 # Create MCH proxy 133 if use_parent_mch: 134 mch_bone = rig.copy_bone(bone, make_derived_name(bone, 'mch', '.parent'), scale=1/3) 135 else: 136 mch_bone = bone 137 138 child = { 139 **self.child_option_table, 140 'rig':rig, 'bone': bone, 'mch_bone': mch_bone, 141 'is_done': False, 'is_configured': False, 142 } 143 self.assign_child_options(child, options) 144 self.child_list.append(child) 145 self.child_map[bone] = child 146 147 148 def amend_child(self, rig, bone, **options): 149 """ 150 Change parameters assigned in a previous build_child call. 151 152 Provided to make it more convenient to change rig behavior by subclassing. 153 """ 154 assert self.generator.stage == 'generate_bones' and not self.frozen 155 child = self.child_map[bone] 156 assert child['rig'] == rig 157 self.assign_child_options(child, options) 158 159 160 def rig_child_now(self, bone): 161 """Create the constraints immediately.""" 162 assert self.generator.stage == 'rig_bones' 163 child = self.child_map[bone] 164 assert not child['is_done'] 165 self.__rig_child(child) 166 167 ############################## 168 # Implementation 169 170 child_option_table = { 171 'extra_parents': None, 172 'prop_bone': None, 'prop_id': None, 'prop_name': None, 'controls': None, 173 'select_parent': None, 'ignore_global': False, 174 'exclude_self': False, 'allow_self': False, 175 'context_rig': None, 'select_tags': None, 176 'no_implicit': False, 'only_selected': False, 177 'ctrl_bone': None, 178 'no_fix_location': False, 'no_fix_rotation': False, 'no_fix_scale': False, 179 'copy_location': None, 'copy_rotation': None, 'copy_scale': None, 180 'inherit_scale': 'AVERAGE', 181 } 182 183 def assign_child_options(self, child, options): 184 if 'context_rig' in options: 185 assert _rig_is_child(child['rig'], options['context_rig']) 186 187 for name, value in options.items(): 188 if name not in self.child_option_table: 189 raise AttributeError('invalid child option: '+name) 190 191 child[name] = value 192 193 def get_rig_parent_candidates(self, rig): 194 candidates = [] 195 196 # Build a list in parent hierarchy order 197 while rig: 198 candidates.append(self.local_parents[id(rig)]) 199 rig = rig.rigify_parent 200 201 candidates.append(self.global_parents) 202 203 return list(chain.from_iterable(reversed(candidates))) 204 205 def generate_bones(self): 206 self.frozen = True 207 self.parent_list = self.global_parents + list(chain.from_iterable(self.local_parents.values())) 208 209 # Link children to parents 210 for child in self.child_list: 211 child_rig = child['context_rig'] or child['rig'] 212 parents = [] 213 214 for parent in self.get_rig_parent_candidates(child_rig): 215 parent_rig = parent['rig'] 216 217 # Exclude injected parents 218 if parent['real_rig'] is not parent_rig: 219 if _rig_is_child(parent_rig, child_rig): 220 continue 221 222 if parent['rig'] is child_rig: 223 if (parent['exclude_self'] and not child['allow_self']) or child['exclude_self']: 224 continue 225 elif parent['is_global'] and not child['ignore_global']: 226 # Can't use parents from own children, even if global (cycle risk) 227 if _rig_is_child(parent_rig, child_rig): 228 continue 229 else: 230 # Required to be a child of the parent's rig 231 if not _rig_is_child(child_rig, parent_rig): 232 continue 233 234 parent['used'] = True 235 parents.append(parent) 236 237 child['parents'] = parents 238 239 # Call lazy creation for parents 240 for parent in self.parent_list: 241 if parent['used']: 242 parent['bone'] = force_lazy(parent['bone']) 243 244 def parent_bones(self): 245 for child in self.child_list: 246 rig = child['rig'] 247 mch = child['mch_bone'] 248 249 # Remove real parent from the child 250 rig.set_bone_parent(mch, None) 251 self.generator.disable_auto_parent(mch) 252 253 # Parent child to the MCH proxy 254 if mch != child['bone']: 255 rig.set_bone_parent(child['bone'], mch, inherit_scale=child['inherit_scale']) 256 257 def configure_bones(self): 258 for child in self.child_list: 259 self.__configure_child(child) 260 261 def __configure_child(self, child): 262 if child['is_configured']: 263 return 264 265 child['is_configured'] = True 266 267 bone = child['bone'] 268 269 # Build the final list of parent bone names 270 parent_map = dict() 271 parent_tags = defaultdict(set) 272 273 for parent in child['parents']: 274 if parent['bone'] not in parent_map: 275 parent_map[parent['bone']] = parent['name'] 276 if parent['tags']: 277 parent_tags[parent['bone']] |= parent['tags'] 278 279 last_main_parent_bone = child['parents'][-1]['bone'] 280 extra_parents = set() 281 282 for parent in force_lazy(child['extra_parents'] or []): 283 if not isinstance(parent, tuple): 284 parent = (parent, None) 285 extra_parents.add(parent[0]) 286 if parent[0] not in parent_map: 287 parent_map[parent[0]] = parent[1] 288 289 for parent in parent_map: 290 if parent in self.child_map: 291 parent_tags[parent] |= {'child'} 292 293 parent_bones = list(parent_map.items()) 294 295 # Find which bone to select 296 select_bone = force_lazy(child['select_parent']) or last_main_parent_bone 297 select_tags = force_lazy(child['select_tags']) or [] 298 299 if child['no_implicit']: 300 assert len(extra_parents) > 0 301 parent_bones = [ item for item in parent_bones if item[0] in extra_parents ] 302 if last_main_parent_bone not in extra_parents: 303 last_main_parent_bone = parent_bones[-1][0] 304 305 for tag in select_tags: 306 tag_set = tag if isinstance(tag, set) else {tag} 307 matching = [ 308 bone for (bone, _) in parent_bones 309 if not tag_set.isdisjoint(parent_tags[bone]) 310 ] 311 if len(matching) > 0: 312 select_bone = matching[-1] 313 break 314 315 if select_bone not in parent_map: 316 print("RIGIFY ERROR: Can't find bone '%s' to select as default parent of '%s'\n" % (select_bone, bone)) 317 select_bone = last_main_parent_bone 318 319 if child['only_selected']: 320 filter_set = { select_bone, *extra_parents } 321 parent_bones = [ item for item in parent_bones if item[0] in filter_set ] 322 323 try: 324 select_index = 1 + next(i for i, (bone, _) in enumerate(parent_bones) if bone == select_bone) 325 except StopIteration: 326 select_index = len(parent_bones) 327 print("RIGIFY ERROR: Invalid default parent '%s' of '%s'\n" % (select_bone, bone)) 328 329 child['parent_bones'] = parent_bones 330 331 # Create the controlling property 332 prop_bone = child['prop_bone'] = force_lazy(child['prop_bone']) or bone 333 prop_name = child['prop_name'] or child['prop_id'] or 'Parent Switch' 334 prop_id = child['prop_id'] = child['prop_id'] or 'parent_switch' 335 336 parent_names = [ parent[1] or strip_prefix(parent[0]) for parent in [(None, 'None'), *parent_bones] ] 337 parent_str = ', '.join([ '%s (%d)' % (name, i) for i, name in enumerate(parent_names) ]) 338 339 ctrl_bone = child['ctrl_bone'] or bone 340 341 self.make_property( 342 prop_bone, prop_id, select_index, 343 min=0, max=len(parent_bones), 344 description='Switch parent of %s: %s' % (ctrl_bone, parent_str) 345 ) 346 347 # Find which channels don't depend on the parent 348 349 no_fix = [ child[n] for n in ['no_fix_location', 'no_fix_rotation', 'no_fix_scale'] ] 350 351 child['copy'] = [ force_lazy(child[n]) for n in ['copy_location', 'copy_rotation', 'copy_scale'] ] 352 353 locks = tuple(bool(nofix or copy) for nofix, copy in zip(no_fix, child['copy'])) 354 355 # Create the script for the property 356 controls = force_lazy(child['controls']) or set([prop_bone, bone]) 357 358 script = self.generator.script 359 panel = script.panel_with_selected_check(child['rig'], controls) 360 361 panel.use_bake_settings() 362 script.add_utilities(SCRIPT_UTILITIES_OP_SWITCH_PARENT) 363 script.register_classes(SCRIPT_REGISTER_OP_SWITCH_PARENT) 364 365 op_props = { 366 'bone': ctrl_bone, 'prop_bone': prop_bone, 'prop_id': prop_id, 367 'parent_names': json.dumps(parent_names), 'locks': locks, 368 } 369 370 row = panel.row(align=True) 371 lsplit = row.split(factor=0.75, align=True) 372 lsplit.operator('pose.rigify_switch_parent_{rig_id}', text=prop_name, icon='DOWNARROW_HLT', properties=op_props) 373 lsplit.custom_prop(prop_bone, prop_id, text='') 374 row.operator('pose.rigify_switch_parent_bake_{rig_id}', text='', icon='ACTION_TWEAK', properties=op_props) 375 376 def rig_bones(self): 377 for child in self.child_list: 378 self.__rig_child(child) 379 380 def __rig_child(self, child): 381 if child['is_done']: 382 return 383 384 child['is_done'] = True 385 386 # Implement via an Armature constraint 387 mch = child['mch_bone'] 388 con = self.make_constraint(mch, 'ARMATURE', name='SWITCH_PARENT') 389 390 prop_var = [(child['prop_bone'], child['prop_id'])] 391 392 for i, (parent, parent_name) in enumerate(child['parent_bones']): 393 tgt = con.targets.new() 394 395 tgt.target = self.obj 396 tgt.subtarget = parent 397 tgt.weight = 0.0 398 399 expr = 'var == %d' % (i+1) 400 self.make_driver(tgt, 'weight', expression=expr, variables=prop_var) 401 402 # Add copy constraints 403 copy = child['copy'] 404 405 if copy[0]: 406 self.make_constraint(mch, 'COPY_LOCATION', copy[0]) 407 if copy[1]: 408 self.make_constraint(mch, 'COPY_ROTATION', copy[1]) 409 if copy[2]: 410 self.make_constraint(mch, 'COPY_SCALE', copy[2]) 411 412 413SCRIPT_REGISTER_OP_SWITCH_PARENT = ['POSE_OT_rigify_switch_parent', 'POSE_OT_rigify_switch_parent_bake'] 414 415SCRIPT_UTILITIES_OP_SWITCH_PARENT = [''' 416################################ 417## Switchable Parent operator ## 418################################ 419 420class RigifySwitchParentBase: 421 bone: StringProperty(name="Control Bone") 422 prop_bone: StringProperty(name="Property Bone") 423 prop_id: StringProperty(name="Property") 424 parent_names: StringProperty(name="Parent Names") 425 locks: bpy.props.BoolVectorProperty(name="Locked", size=3, default=[False,False,False]) 426 427 parent_items = [('0','None','None')] 428 429 selected: bpy.props.EnumProperty( 430 name='Selected Parent', 431 items=lambda s,c: RigifySwitchParentBase.parent_items 432 ) 433 434 def save_frame_state(self, context, obj): 435 return get_transform_matrix(obj, self.bone, with_constraints=False) 436 437 def apply_frame_state(self, context, obj, old_matrix): 438 # Change the parent 439 set_custom_property_value( 440 obj, self.prop_bone, self.prop_id, int(self.selected), 441 keyflags=self.keyflags_switch 442 ) 443 444 context.view_layer.update() 445 446 # Set the transforms to restore position 447 set_transform_from_matrix( 448 obj, self.bone, old_matrix, keyflags=self.keyflags, 449 no_loc=self.locks[0], no_rot=self.locks[1], no_scale=self.locks[2] 450 ) 451 452 def init_invoke(self, context): 453 pose = context.active_object.pose 454 455 if (not pose or not self.parent_names 456 or self.bone not in pose.bones 457 or self.prop_bone not in pose.bones 458 or self.prop_id not in pose.bones[self.prop_bone]): 459 self.report({'ERROR'}, "Invalid parameters") 460 return {'CANCELLED'} 461 462 parents = json.loads(self.parent_names) 463 pitems = [(str(i), name, name) for i, name in enumerate(parents)] 464 465 RigifySwitchParentBase.parent_items = pitems 466 467 self.selected = str(pose.bones[self.prop_bone][self.prop_id]) 468 469 470class POSE_OT_rigify_switch_parent(RigifySwitchParentBase, RigifySingleUpdateMixin, bpy.types.Operator): 471 bl_idname = "pose.rigify_switch_parent_" + rig_id 472 bl_label = "Switch Parent (Keep Transform)" 473 bl_options = {'REGISTER', 'UNDO', 'INTERNAL'} 474 bl_description = "Switch parent, preserving the bone position and orientation" 475 476 def draw(self, _context): 477 col = self.layout.column() 478 col.prop(self, 'selected', expand=True) 479 480 481class POSE_OT_rigify_switch_parent_bake(RigifySwitchParentBase, RigifyBakeKeyframesMixin, bpy.types.Operator): 482 bl_idname = "pose.rigify_switch_parent_bake_" + rig_id 483 bl_label = "Apply Switch Parent To Keyframes" 484 bl_description = "Switch parent over a frame range, adjusting keys to preserve the bone position and orientation" 485 486 def execute_scan_curves(self, context, obj): 487 return self.bake_add_bone_frames(self.bone, transform_props_with_locks(*self.locks)) 488 489 def execute_before_apply(self, context, obj, range, range_raw): 490 self.bake_replace_custom_prop_keys_constant(self.prop_bone, self.prop_id, int(self.selected)) 491 492 def draw(self, context): 493 self.layout.prop(self, 'selected', text='') 494'''] 495