rigify/utils/switch_parent.py

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