rigify/base_generate.py

   1 # SPDX-License-Identifier: GPL-2.0-or-later
   2
   3 import bpy
   4 import sys
   5 import traceback
   6 import collections
   7
   8 from .utils.errors import MetarigError, RaiseErrorMixin
   9 from .utils.naming import random_id
  10 from .utils.metaclass import SingletonPluginMetaclass
  11 from .utils.rig import list_bone_names_depth_first_sorted, get_rigify_type
  12 from .utils.misc import clone_parameters, assign_parameters
  13
  14 from . import base_rig
  15
  16 from itertools import count
  17
  18 #=============================================
  19 # Generator Plugin
  20 #=============================================
  21
  22
  23 class GeneratorPlugin(base_rig.GenerateCallbackHost, metaclass=SingletonPluginMetaclass):
  24     """
  25     Base class for generator plugins.
  26
  27     Generator plugins are per-Generator singleton utility
  28     classes that receive the same stage callbacks as rigs.
  29
  30     Useful for building entities shared by multiple rigs
  31     (e.g. the python script), or for making fire-and-forget
  32     utilities that actually require multiple stages to
  33     complete.
  34
  35     This will create only one instance per set of args:
  36
  37       instance = PluginClass(generator, ...init args)
  38     """
  39
  40     priority = 0
  41
  42     def __init__(self, generator):
  43         self.generator = generator
  44         self.obj = generator.obj
  45
  46     def register_new_bone(self, new_name, old_name=None):
  47         self.generator.bone_owners[new_name] = None
  48         if old_name:
  49             self.generator.derived_bones[old_name].add(new_name)
  50
  51
  52 #=============================================
  53 # Rig Substitution Mechanism
  54 #=============================================
  55
  56
  57 class SubstitutionRig(RaiseErrorMixin):
  58     """A proxy rig that replaces itself with one or more different rigs."""
  59
  60     def __init__(self, generator, pose_bone):
  61         self.generator = generator
  62
  63         self.obj = generator.obj
  64         self.base_bone = pose_bone.name
  65         self.params = pose_bone.rigify_parameters
  66         self.params_copy = clone_parameters(self.params)
  67
  68     def substitute(self):
  69         # return [rig1, rig2...]
  70         raise NotImplementedException()
  71
  72     # Utility methods
  73     def register_new_bone(self, new_name, old_name=None):
  74         pass
  75
  76     def get_params(self, bone_name):
  77         return self.obj.pose.bones[bone_name].rigify_parameters
  78
  79     def assign_params(self, bone_name, param_dict=None, **params):
  80         assign_parameters(self.get_params(bone_name), param_dict, **params)
  81
  82     def instantiate_rig(self, rig_class, bone_name):
  83         if isinstance(rig_class, str):
  84             rig_class = self.generator.find_rig_class(rig_class)
  85
  86         return self.generator.instantiate_rig(rig_class, self.obj.pose.bones[bone_name])
  87
  88
  89 #=============================================
  90 # Legacy Rig Wrapper
  91 #=============================================
  92
  93
  94 class LegacyRig(base_rig.BaseRig):
  95     """Wrapper around legacy style rigs without a common base class"""
  96
  97     def __init__(self, generator, pose_bone, wrapped_class):
  98         self.wrapped_rig = None
  99         self.wrapped_class = wrapped_class
 100
 101         super().__init__(generator, pose_bone)
 102
 103     def find_org_bones(self, pose_bone):
 104         bone_name = pose_bone.name
 105
 106         if not self.wrapped_rig:
 107             self.wrapped_rig = self.wrapped_class(self.obj, self.base_bone, self.params)
 108
 109             # Switch back to OBJECT mode if the rig changed it
 110             if self.obj.mode != 'OBJECT':
 111                 bpy.ops.object.mode_set(mode='OBJECT')
 112
 113         # Try to extract the main list of bones - old rigs often have it.
 114         # This is not actually strictly necessary, so failing is OK.
 115         if hasattr(self.wrapped_rig, 'org_bones'):
 116             bones = self.wrapped_rig.org_bones
 117             if isinstance(bones, list):
 118                 return bones
 119
 120         return [bone_name]
 121
 122     def generate_bones(self):
 123         # Inject references into the rig if it won't cause conflict
 124         if not hasattr(self.wrapped_rig, 'rigify_generator'):
 125             self.wrapped_rig.rigify_generator = self.generator
 126         if not hasattr(self.wrapped_rig, 'rigify_wrapper'):
 127             self.wrapped_rig.rigify_wrapper = self
 128
 129         # Old rigs only have one generate method, so call it from
 130         # generate_bones, which is the only stage allowed to add bones.
 131         scripts = self.wrapped_rig.generate()
 132
 133         # Switch back to EDIT mode if the rig changed it
 134         if self.obj.mode != 'EDIT':
 135             bpy.ops.object.mode_set(mode='EDIT')
 136
 137         if isinstance(scripts, dict):
 138             if 'script' in scripts:
 139                 self.script.add_panel_code(scripts['script'])
 140             if 'imports' in scripts:
 141                 self.script.add_imports(scripts['imports'])
 142             if 'utilities' in scripts:
 143                 self.script.add_utilities(scripts['utilities'])
 144             if 'register' in scripts:
 145                 self.script.register_classes(scripts['register'])
 146             if 'register_drivers' in scripts:
 147                 self.script.register_driver_functions(scripts['register_drivers'])
 148             if 'register_props' in scripts:
 149                 for prop, val in scripts['register_props']:
 150                     self.script.register_property(prop, val)
 151             if 'noparent_bones' in scripts:
 152                 for bone_name in scripts['noparent_bones']:
 153                     self.generator.disable_auto_parent(bone_name)
 154         elif scripts is not None:
 155             self.script.add_panel_code([scripts[0]])
 156
 157     def finalize(self):
 158         if hasattr(self.wrapped_rig, 'glue'):
 159             self.wrapped_rig.glue()
 160
 161             # Switch back to OBJECT mode if the rig changed it
 162             if self.obj.mode != 'OBJECT':
 163                 bpy.ops.object.mode_set(mode='OBJECT')
 164
 165
 166 #=============================================
 167 # Base Generate Engine
 168 #=============================================
 169
 170
 171 class BaseGenerator:
 172     """Base class for the main generator object. Contains rig and plugin management code."""
 173
 174     instance = None
 175
 176     def __init__(self, context, metarig):
 177         self.context = context
 178         self.scene = context.scene
 179         self.view_layer = context.view_layer
 180         self.layer_collection = context.layer_collection
 181         self.collection = self.layer_collection.collection
 182         self.metarig = metarig
 183         self.obj = None
 184
 185         # List of all rig instances
 186         self.rig_list = []
 187         # List of rigs that don't have a parent
 188         self.root_rigs = []
 189         # Map from bone names to their rigs
 190         self.bone_owners = {}
 191         self.derived_bones = collections.defaultdict(set)
 192
 193         # Set of plugins
 194         self.plugin_list = []
 195         self.plugin_map = {}
 196
 197         # Current execution stage so plugins could check they are used correctly
 198         self.stage = None
 199
 200         # Set of bones that should be left without parent
 201         self.noparent_bones = set()
 202
 203         # Table of layer priorities for defining bone groups
 204         self.layer_group_priorities = collections.defaultdict(dict)
 205
 206         # Random string with time appended so that
 207         # different rigs don't collide id's
 208         self.rig_id = random_id(16)
 209
 210         # Table of renamed ORG bones
 211         self.org_rename_table = dict()
 212
 213
 214     def disable_auto_parent(self, bone_name):
 215         """Prevent automatically parenting the bone to root if parentless."""
 216         self.noparent_bones.add(bone_name)
 217
 218
 219     def find_derived_bones(self, bone_name, *, by_owner=False, recursive=True):
 220         """Find which bones were copied from the specified one."""
 221         if by_owner:
 222             owner = self.bone_owners.get(bone_name, None)
 223             if not owner:
 224                 return {}
 225
 226             table = owner.rigify_derived_bones
 227         else:
 228             table = self.derived_bones
 229
 230         if recursive:
 231             result = set()
 232
 233             def rec(name):
 234                 for child in table.get(name, {}):
 235                     result.add(child)
 236                     rec(child)
 237
 238             rec(bone_name)
 239
 240             return result
 241         else:
 242             return set(table.get(bone_name, {}))
 243
 244
 245     def set_layer_group_priority(self, bone_name, layers, priority):
 246         for i, val in enumerate(layers):
 247             if val:
 248                 self.layer_group_priorities[bone_name][i] = priority
 249
 250
 251     def rename_org_bone(self, old_name, new_name):
 252         assert self.stage == 'instantiate'
 253         assert old_name == self.org_rename_table.get(old_name, None)
 254         assert old_name not in self.bone_owners
 255
 256         bone = self.obj.data.bones[old_name]
 257
 258         bone.name = new_name
 259         new_name = bone.name
 260
 261         self.org_rename_table[old_name] = new_name
 262         return new_name
 263
 264
 265     def __run_object_stage(self, method_name):
 266         assert(self.context.active_object == self.obj)
 267         assert(self.obj.mode == 'OBJECT')
 268         num_bones = len(self.obj.data.bones)
 269
 270         self.stage = method_name
 271
 272         for rig in self.rig_list:
 273             rig.rigify_invoke_stage(method_name)
 274
 275             assert(self.context.active_object == self.obj)
 276             assert(self.obj.mode == 'OBJECT')
 277             assert(num_bones == len(self.obj.data.bones))
 278
 279         # Allow plugins to be added to the end of the list on the fly
 280         for i in count(0):
 281             if i >= len(self.plugin_list):
 282                 break
 283
 284             self.plugin_list[i].rigify_invoke_stage(method_name)
 285
 286             assert(self.context.active_object == self.obj)
 287             assert(self.obj.mode == 'OBJECT')
 288             assert(num_bones == len(self.obj.data.bones))
 289
 290
 291     def __run_edit_stage(self, method_name):
 292         assert(self.context.active_object == self.obj)
 293         assert(self.obj.mode == 'EDIT')
 294         num_bones = len(self.obj.data.edit_bones)
 295
 296         self.stage = method_name
 297
 298         for rig in self.rig_list:
 299             rig.rigify_invoke_stage(method_name)
 300
 301             assert(self.context.active_object == self.obj)
 302             assert(self.obj.mode == 'EDIT')
 303             assert(num_bones == len(self.obj.data.edit_bones))
 304
 305         # Allow plugins to be added to the end of the list on the fly
 306         for i in count(0):
 307             if i >= len(self.plugin_list):
 308                 break
 309
 310             self.plugin_list[i].rigify_invoke_stage(method_name)
 311
 312             assert(self.context.active_object == self.obj)
 313             assert(self.obj.mode == 'EDIT')
 314             assert(num_bones == len(self.obj.data.edit_bones))
 315
 316
 317     def invoke_initialize(self):
 318         self.__run_object_stage('initialize')
 319
 320
 321     def invoke_prepare_bones(self):
 322         self.__run_edit_stage('prepare_bones')
 323
 324
 325     def __auto_register_bones(self, bones, rig, plugin=None):
 326         """Find bones just added and not registered by this rig."""
 327         for bone in bones:
 328             name = bone.name
 329             if name not in self.bone_owners:
 330                 self.bone_owners[name] = rig
 331                 if rig:
 332                     rig.rigify_new_bones[name] = None
 333
 334                     if not isinstance(rig, LegacyRig):
 335                         print("WARNING: rig %s didn't register bone %s\n" % (self.describe_rig(rig), name))
 336                 else:
 337                     print("WARNING: plugin %s didn't register bone %s\n" % (plugin, name))
 338
 339
 340     def invoke_generate_bones(self):
 341         assert(self.context.active_object == self.obj)
 342         assert(self.obj.mode == 'EDIT')
 343
 344         self.stage = 'generate_bones'
 345
 346         for rig in self.rig_list:
 347             rig.rigify_invoke_stage('generate_bones')
 348
 349             assert(self.context.active_object == self.obj)
 350             assert(self.obj.mode == 'EDIT')
 351
 352             self.__auto_register_bones(self.obj.data.edit_bones, rig)
 353
 354         # Allow plugins to be added to the end of the list on the fly
 355         for i in count(0):
 356             if i >= len(self.plugin_list):
 357                 break
 358
 359             self.plugin_list[i].rigify_invoke_stage('generate_bones')
 360
 361             assert(self.context.active_object == self.obj)
 362             assert(self.obj.mode == 'EDIT')
 363
 364             self.__auto_register_bones(self.obj.data.edit_bones, None, plugin=self.plugin_list[i])
 365
 366
 367     def invoke_parent_bones(self):
 368         self.__run_edit_stage('parent_bones')
 369
 370
 371     def invoke_configure_bones(self):
 372         self.__run_object_stage('configure_bones')
 373
 374
 375     def invoke_preapply_bones(self):
 376         self.__run_object_stage('preapply_bones')
 377
 378
 379     def invoke_apply_bones(self):
 380         self.__run_edit_stage('apply_bones')
 381
 382
 383     def invoke_rig_bones(self):
 384         self.__run_object_stage('rig_bones')
 385
 386
 387     def invoke_generate_widgets(self):
 388         self.__run_object_stage('generate_widgets')
 389
 390
 391     def invoke_finalize(self):
 392         self.__run_object_stage('finalize')
 393
 394
 395     def instantiate_rig(self, rig_class, pose_bone):
 396         assert not issubclass(rig_class, SubstitutionRig)
 397
 398         if issubclass(rig_class, base_rig.BaseRig):
 399             return rig_class(self, pose_bone)
 400         else:
 401             return LegacyRig(self, pose_bone, rig_class)
 402
 403
 404     def instantiate_rig_by_type(self, rig_type, pose_bone):
 405         return self.instantiate_rig(self.find_rig_class(rig_type), pose_bone)
 406
 407
 408     def describe_rig(self, rig):
 409         base_bone = rig.base_bone
 410
 411         if isinstance(rig, LegacyRig):
 412             rig = rig.wrapped_rig
 413
 414         return "%s (%s)" % (rig.__class__, base_bone)
 415
 416
 417     def __create_rigs(self, bone_name, halt_on_missing):
 418         """Recursively walk bones and create rig instances."""
 419
 420         pose_bone = self.obj.pose.bones[bone_name]
 421
 422         rig_type = get_rigify_type(pose_bone)
 423
 424         if rig_type != "":
 425             try:
 426                 rig_class = self.find_rig_class(rig_type)
 427
 428                 if issubclass(rig_class, SubstitutionRig):
 429                     rigs = rig_class(self, pose_bone).substitute()
 430                 else:
 431                     rigs = [self.instantiate_rig(rig_class, pose_bone)]
 432
 433                 assert(self.context.active_object == self.obj)
 434                 assert(self.obj.mode == 'OBJECT')
 435
 436                 for rig in rigs:
 437                     self.rig_list.append(rig)
 438
 439                     for org_name in rig.rigify_org_bones:
 440                         if org_name in self.bone_owners:
 441                             old_rig = self.describe_rig(self.bone_owners[org_name])
 442                             new_rig = self.describe_rig(rig)
 443                             print("CONFLICT: bone %s is claimed by rigs %s and %s\n" % (org_name, old_rig, new_rig))
 444
 445                         self.bone_owners[org_name] = rig
 446
 447             except ImportError:
 448                 message = "Rig Type Missing: python module for type '%s' not found (bone: %s)" % (rig_type, bone_name)
 449                 if halt_on_missing:
 450                     raise MetarigError(message)
 451                 else:
 452                     print(message)
 453                     print('print_exc():')
 454                     traceback.print_exc(file=sys.stdout)
 455
 456
 457     def __build_rig_tree_rec(self, bone, current_rig, handled):
 458         """Recursively walk bones and connect rig instances into a tree."""
 459
 460         rig = self.bone_owners.get(bone.name)
 461
 462         if rig:
 463             if rig is current_rig:
 464                 pass
 465
 466             elif rig not in handled:
 467                 rig.rigify_parent = current_rig
 468
 469                 if current_rig:
 470                     current_rig.rigify_children.append(rig)
 471                 else:
 472                     self.root_rigs.append(rig)
 473
 474                 handled[rig] = bone.name
 475
 476             elif rig.rigify_parent is not current_rig:
 477                 raise MetarigError("CONFLICT: bone %s owned by rig %s has different parent rig from %s\n" %
 478                                    (bone.name, rig.base_bone, handled[rig]))
 479
 480             current_rig = rig
 481         else:
 482             if current_rig:
 483                 current_rig.rigify_child_bones.add(bone.name)
 484
 485             self.bone_owners[bone.name] = current_rig
 486
 487         for child in bone.children:
 488             self.__build_rig_tree_rec(child, current_rig, handled)
 489
 490
 491     def instantiate_rig_tree(self, halt_on_missing=False):
 492         """Create rig instances and connect them into a tree."""
 493
 494         assert(self.context.active_object == self.obj)
 495         assert(self.obj.mode == 'OBJECT')
 496
 497         self.stage = 'instantiate'
 498
 499         # Compute the list of bones
 500         bone_list = list_bone_names_depth_first_sorted(self.obj)
 501
 502         self.org_rename_table = {n: n for n in bone_list}
 503
 504         # Construct the rig instances
 505         for name in bone_list:
 506             self.__create_rigs(self.org_rename_table[name], halt_on_missing)
 507
 508         # Connect rigs and bones into a tree
 509         handled = {}
 510
 511         for bone in self.obj.data.bones:
 512             if bone.parent is None:
 513                 self.__build_rig_tree_rec(bone, None, handled)