rigify/utils/switch_parent.py

   1 import bpy
   2
   3 import re
   4 import itertools
   5 import bisect
   6 import json
   7
   8 from .errors import MetarigError
   9 from .naming import strip_prefix, make_derived_name
  10 from .mechanism import MechanismUtilityMixin
  11 from .misc import map_list, map_apply, force_lazy
  12
  13 from ..base_rig import *
  14 from ..base_generate import GeneratorPlugin
  15
  16 from collections import defaultdict
  17 from itertools import count, repeat, chain
  18
  19
  20 def _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
  33 class 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(
 389             mch, 'ARMATURE', name='SWITCH_PARENT',
 390             targets=[ (parent, 0.0) for parent, _ in child['parent_bones'] ]
 391         )
 392
 393         prop_var = [(child['prop_bone'], child['prop_id'])]
 394
 395         for i, (parent, parent_name) in enumerate(child['parent_bones']):
 396             expr = 'var == %d' % (i+1)
 397             self.make_driver(con.targets[i], 'weight', expression=expr, variables=prop_var)
 398
 399         # Add copy constraints
 400         copy = child['copy']
 401
 402         if copy[0]:
 403             self.make_constraint(mch, 'COPY_LOCATION', copy[0])
 404         if copy[1]:
 405             self.make_constraint(mch, 'COPY_ROTATION', copy[1])
 406         if copy[2]:
 407             self.make_constraint(mch, 'COPY_SCALE', copy[2])
 408
 409
 410 SCRIPT_REGISTER_OP_SWITCH_PARENT = ['POSE_OT_rigify_switch_parent', 'POSE_OT_rigify_switch_parent_bake']
 411
 412 SCRIPT_UTILITIES_OP_SWITCH_PARENT = ['''
 413 ################################
 414 ## Switchable Parent operator ##
 415 ################################
 416
 417 class RigifySwitchParentBase:
 418     bone:         StringProperty(name="Control Bone")
 419     prop_bone:    StringProperty(name="Property Bone")
 420     prop_id:      StringProperty(name="Property")
 421     parent_names: StringProperty(name="Parent Names")
 422     locks:        bpy.props.BoolVectorProperty(name="Locked", size=3, default=[False,False,False])
 423
 424     parent_items = [('0','None','None')]
 425
 426     selected: bpy.props.EnumProperty(
 427         name='Selected Parent',
 428         items=lambda s,c: RigifySwitchParentBase.parent_items
 429     )
 430
 431     def save_frame_state(self, context, obj):
 432         return get_transform_matrix(obj, self.bone, with_constraints=False)
 433
 434     def apply_frame_state(self, context, obj, old_matrix):
 435         # Change the parent
 436         set_custom_property_value(
 437             obj, self.prop_bone, self.prop_id, int(self.selected),
 438             keyflags=self.keyflags_switch
 439         )
 440
 441         context.view_layer.update()
 442
 443         # Set the transforms to restore position
 444         set_transform_from_matrix(
 445             obj, self.bone, old_matrix, keyflags=self.keyflags,
 446             no_loc=self.locks[0], no_rot=self.locks[1], no_scale=self.locks[2]
 447         )
 448
 449     def init_invoke(self, context):
 450         pose = context.active_object.pose
 451
 452         if (not pose or not self.parent_names
 453             or self.bone not in pose.bones
 454             or self.prop_bone not in pose.bones
 455             or self.prop_id not in pose.bones[self.prop_bone]):
 456             self.report({'ERROR'}, "Invalid parameters")
 457             return {'CANCELLED'}
 458
 459         parents = json.loads(self.parent_names)
 460         pitems = [(str(i), name, name) for i, name in enumerate(parents)]
 461
 462         RigifySwitchParentBase.parent_items = pitems
 463
 464         self.selected = str(pose.bones[self.prop_bone][self.prop_id])
 465
 466
 467 class POSE_OT_rigify_switch_parent(RigifySwitchParentBase, RigifySingleUpdateMixin, bpy.types.Operator):
 468     bl_idname = "pose.rigify_switch_parent_" + rig_id
 469     bl_label = "Switch Parent (Keep Transform)"
 470     bl_options = {'REGISTER', 'UNDO', 'INTERNAL'}
 471     bl_description = "Switch parent, preserving the bone position and orientation"
 472
 473     def draw(self, _context):
 474         col = self.layout.column()
 475         col.prop(self, 'selected', expand=True)
 476
 477
 478 class POSE_OT_rigify_switch_parent_bake(RigifySwitchParentBase, RigifyBakeKeyframesMixin, bpy.types.Operator):
 479     bl_idname = "pose.rigify_switch_parent_bake_" + rig_id
 480     bl_label = "Apply Switch Parent To Keyframes"
 481     bl_description = "Switch parent over a frame range, adjusting keys to preserve the bone position and orientation"
 482
 483     def execute_scan_curves(self, context, obj):
 484         return self.bake_add_bone_frames(self.bone, transform_props_with_locks(*self.locks))
 485
 486     def execute_before_apply(self, context, obj, range, range_raw):
 487         self.bake_replace_custom_prop_keys_constant(self.prop_bone, self.prop_id, int(self.selected))
 488
 489     def draw(self, context):
 490         self.layout.prop(self, 'selected', text='')
 491 ''']