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 typing import Optional, TYPE_CHECKING, Collection, List
   9 from bpy.types import PoseBone, Bone
  10
  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
  16
  17 from . import base_rig
  18
  19 from itertools import count
  20
  21 if TYPE_CHECKING:
  22     from .rig_ui_template import ScriptGenerator
  23
  24
  25 ##############################################
  26 # Generator Plugin
  27 ##############################################
  28
  29
  30 class GeneratorPlugin(base_rig.GenerateCallbackHost, metaclass=SingletonPluginMetaclass):
  31     """
  32     Base class for generator plugins.
  33
  34     Generator plugins are per-Generator singleton utility
  35     classes that receive the same stage callbacks as rigs.
  36
  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
  40     complete.
  41
  42     This will create only one instance per set of args:
  43
  44       instance = PluginClass(generator, ...init args)
  45     """
  46
  47     priority = 0
  48
  49     def __init__(self, generator: 'BaseGenerator'):
  50         self.generator = generator
  51         self.obj = generator.obj
  52
  53     def register_new_bone(self, new_name: str, old_name: Optional[str] = None):
  54         self.generator.bone_owners[new_name] = None
  55         if old_name:
  56             self.generator.derived_bones[old_name].add(new_name)
  57
  58
  59 ##############################################
  60 # Rig Substitution Mechanism
  61 ##############################################
  62
  63
  64 class SubstitutionRig(RaiseErrorMixin):
  65     """A proxy rig that replaces itself with one or more different rigs."""
  66
  67     def __init__(self, generator: 'BaseGenerator', pose_bone: PoseBone):
  68         self.generator = generator
  69
  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)
  74
  75     def substitute(self):
  76         # return [rig1, rig2...]
  77         raise NotImplementedError
  78
  79     # Utility methods
  80     def register_new_bone(self, new_name: str, old_name: Optional[str] = None):
  81         pass
  82
  83     def get_params(self, bone_name: str):
  84         return get_rigify_params(self.obj.pose.bones[bone_name])
  85
  86     def assign_params(self, bone_name: str, param_dict=None, **params):
  87         assign_parameters(self.get_params(bone_name), param_dict, **params)
  88
  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)
  92
  93         return self.generator.instantiate_rig(rig_class, self.obj.pose.bones[bone_name])
  94
  95
  96 ##############################################
  97 # Legacy Rig Wrapper
  98 ##############################################
  99
 100
 101 class LegacyRig(base_rig.BaseRig):
 102     """Wrapper around legacy style rigs without a common base class"""
 103
 104     def __init__(self, generator: 'BaseGenerator', pose_bone: PoseBone, wrapped_class: type):
 105         self.wrapped_rig = None
 106         self.wrapped_class = wrapped_class
 107
 108         super().__init__(generator, pose_bone)
 109
 110     def find_org_bones(self, pose_bone: PoseBone):
 111         bone_name = pose_bone.name
 112
 113         if not self.wrapped_rig:
 114             self.wrapped_rig = self.wrapped_class(self.obj, self.base_bone, self.params)
 115
 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')
 119
 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):
 125                 return bones
 126
 127         return [bone_name]
 128
 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
 135
 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()
 139
 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')
 143
 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]])
 163
 164     def finalize(self):
 165         if hasattr(self.wrapped_rig, 'glue'):
 166             self.wrapped_rig.glue()
 167
 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')
 171
 172
 173 ##############################################
 174 # Base Generate Engine
 175 ##############################################
 176
 177
 178 class BaseGenerator:
 179     """Base class for the main generator object. Contains rig and plugin management code."""
 180
 181     instance: Optional['BaseGenerator'] = None  # static
 182
 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
 188
 189     metarig: ArmatureObject
 190     obj: ArmatureObject
 191
 192     script: 'ScriptGenerator'
 193
 194     rig_list: List[base_rig.BaseRig]
 195     root_rigs: List[base_rig.BaseRig]
 196
 197     bone_owners: dict[str, Optional[base_rig.BaseRig]]
 198     derived_bones: dict[str, set[str]]
 199
 200     stage: Optional[str]
 201     rig_id: str
 202
 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]
 208
 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
 216
 217         # List of all rig instances
 218         self.rig_list = []
 219         # List of rigs that don't have a parent
 220         self.root_rigs = []
 221         # Map from bone names to their rigs
 222         self.bone_owners = {}
 223         self.derived_bones = collections.defaultdict(set)
 224
 225         # Set of plugins
 226         self.plugin_list = []
 227         self.plugin_map = {}
 228
 229         # Current execution stage so plugins could check they are used correctly
 230         self.stage = None
 231
 232         # Set of bones that should be left without parent
 233         self.noparent_bones = set()
 234
 235         # Table of layer priorities for defining bone groups
 236         self.layer_group_priorities = collections.defaultdict(dict)
 237
 238         # Random string with time appended so that
 239         # different rigs don't collide id's
 240         self.rig_id = random_id(16)
 241
 242         # Table of renamed ORG bones
 243         self.org_rename_table = dict()
 244
 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)
 248
 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."""
 251         if by_owner:
 252             owner = self.bone_owners.get(bone_name, None)
 253             if not owner:
 254                 return set()
 255
 256             table = owner.rigify_derived_bones
 257         else:
 258             table = self.derived_bones
 259
 260         if recursive:
 261             result = set()
 262
 263             def rec(name):
 264                 for child in table.get(name, []):
 265                     result.add(child)
 266                     rec(child)
 267
 268             rec(bone_name)
 269
 270             return result
 271         else:
 272             return set(table.get(bone_name, []))
 273
 274     def set_layer_group_priority(self, bone_name: str,
 275                                  layers: Collection[bool], priority: float):
 276         for i, val in enumerate(layers):
 277             if val:
 278                 self.layer_group_priorities[bone_name][i] = priority
 279
 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
 284
 285         bone = self.obj.data.bones[old_name]
 286
 287         bone.name = new_name
 288         new_name = bone.name
 289
 290         self.org_rename_table[old_name] = new_name
 291         return new_name
 292
 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)
 298
 299         self.stage = method_name
 300
 301         for rig in self.rig_list:
 302             rig.rigify_invoke_stage(method_name)
 303
 304             assert(self.context.active_object == self.obj)
 305             assert(self.obj.mode == 'OBJECT')
 306             assert(num_bones == len(self.obj.data.bones))
 307
 308         # Allow plugins to be added to the end of the list on the fly
 309         for i in count(0):
 310             if i >= len(self.plugin_list):
 311                 break
 312
 313             self.plugin_list[i].rigify_invoke_stage(method_name)
 314
 315             assert(self.context.active_object == self.obj)
 316             assert(self.obj.mode == 'OBJECT')
 317             assert(num_bones == len(self.obj.data.bones))
 318
 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)
 324
 325         self.stage = method_name
 326
 327         for rig in self.rig_list:
 328             rig.rigify_invoke_stage(method_name)
 329
 330             assert(self.context.active_object == self.obj)
 331             assert(self.obj.mode == 'EDIT')
 332             assert(num_bones == len(self.obj.data.edit_bones))
 333
 334         # Allow plugins to be added to the end of the list on the fly
 335         for i in count(0):
 336             if i >= len(self.plugin_list):
 337                 break
 338
 339             self.plugin_list[i].rigify_invoke_stage(method_name)
 340
 341             assert(self.context.active_object == self.obj)
 342             assert(self.obj.mode == 'EDIT')
 343             assert(num_bones == len(self.obj.data.edit_bones))
 344
 345     def invoke_initialize(self):
 346         self.__run_object_stage('initialize')
 347
 348     def invoke_prepare_bones(self):
 349         self.__run_edit_stage('prepare_bones')
 350
 351     def __auto_register_bones(self, bones, rig, plugin=None):
 352         """Find bones just added and not registered by this rig."""
 353         for bone in bones:
 354             name = bone.name
 355             if name not in self.bone_owners:
 356                 self.bone_owners[name] = rig
 357                 if rig:
 358                     rig.rigify_new_bones[name] = None
 359
 360                     if not isinstance(rig, LegacyRig):
 361                         print(f"WARNING: rig {self.describe_rig(rig)} "
 362                               f"didn't register bone {name}\n")
 363                 else:
 364                     print(f"WARNING: plugin {plugin} didn't register bone {name}\n")
 365
 366     def invoke_generate_bones(self):
 367         assert(self.context.active_object == self.obj)
 368         assert(self.obj.mode == 'EDIT')
 369
 370         self.stage = 'generate_bones'
 371
 372         for rig in self.rig_list:
 373             rig.rigify_invoke_stage('generate_bones')
 374
 375             assert(self.context.active_object == self.obj)
 376             assert(self.obj.mode == 'EDIT')
 377
 378             self.__auto_register_bones(self.obj.data.edit_bones, rig)
 379
 380         # Allow plugins to be added to the end of the list on the fly
 381         for i in count(0):
 382             if i >= len(self.plugin_list):
 383                 break
 384
 385             self.plugin_list[i].rigify_invoke_stage('generate_bones')
 386
 387             assert(self.context.active_object == self.obj)
 388             assert(self.obj.mode == 'EDIT')
 389
 390             self.__auto_register_bones(self.obj.data.edit_bones, None, plugin=self.plugin_list[i])
 391
 392     def invoke_parent_bones(self):
 393         self.__run_edit_stage('parent_bones')
 394
 395     def invoke_configure_bones(self):
 396         self.__run_object_stage('configure_bones')
 397
 398     def invoke_preapply_bones(self):
 399         self.__run_object_stage('preapply_bones')
 400
 401     def invoke_apply_bones(self):
 402         self.__run_edit_stage('apply_bones')
 403
 404     def invoke_rig_bones(self):
 405         self.__run_object_stage('rig_bones')
 406
 407     def invoke_generate_widgets(self):
 408         self.__run_object_stage('generate_widgets')
 409
 410     def invoke_finalize(self):
 411         self.__run_object_stage('finalize')
 412
 413     def instantiate_rig(self, rig_class: type, pose_bone: PoseBone) -> base_rig.BaseRig:
 414         assert not issubclass(rig_class, SubstitutionRig)
 415
 416         if issubclass(rig_class, base_rig.BaseRig):
 417             return rig_class(self, pose_bone)
 418         else:
 419             return LegacyRig(self, pose_bone, rig_class)
 420
 421     def find_rig_class(self, rig_type: str) -> type:
 422         raise NotImplementedError
 423
 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)
 426
 427     # noinspection PyMethodMayBeStatic
 428     def describe_rig(self, rig: base_rig.BaseRig) -> str:
 429         base_bone = rig.base_bone
 430
 431         if isinstance(rig, LegacyRig):
 432             rig = rig.wrapped_rig
 433
 434         return "%s (%s)" % (rig.__class__, base_bone)
 435
 436     def __create_rigs(self, bone_name, halt_on_missing):
 437         """Recursively walk bones and create rig instances."""
 438
 439         pose_bone = self.obj.pose.bones[bone_name]
 440
 441         rig_type = get_rigify_type(pose_bone)
 442
 443         if rig_type != "":
 444             try:
 445                 rig_class = self.find_rig_class(rig_type)
 446
 447                 if issubclass(rig_class, SubstitutionRig):
 448                     rigs = rig_class(self, pose_bone).substitute()
 449                 else:
 450                     rigs = [self.instantiate_rig(rig_class, pose_bone)]
 451
 452                 assert(self.context.active_object == self.obj)
 453                 assert(self.obj.mode == 'OBJECT')
 454
 455                 for rig in rigs:
 456                     self.rig_list.append(rig)
 457
 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")
 464
 465                         self.bone_owners[org_name] = rig
 466
 467             except ImportError:
 468                 message = f"Rig Type Missing: python module for type '{rig_type}' "\
 469                           f"not found (bone: {bone_name})"
 470                 if halt_on_missing:
 471                     raise MetarigError(message)
 472                 else:
 473                     print(message)
 474                     print('print_exc():')
 475                     traceback.print_exc(file=sys.stdout)
 476
 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."""
 480
 481         rig = self.bone_owners.get(bone.name)
 482
 483         if rig:
 484             if rig is current_rig:
 485                 pass
 486
 487             elif rig not in handled:
 488                 rig.rigify_parent = current_rig
 489
 490                 if current_rig:
 491                     current_rig.rigify_children.append(rig)
 492                 else:
 493                     self.root_rigs.append(rig)
 494
 495                 handled[rig] = bone.name
 496
 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]}")
 500
 501             current_rig = rig
 502         else:
 503             if current_rig:
 504                 current_rig.rigify_child_bones.add(bone.name)
 505
 506             self.bone_owners[bone.name] = current_rig
 507
 508         for child in bone.children:
 509             self.__build_rig_tree_rec(child, current_rig, handled)
 510
 511     def instantiate_rig_tree(self, halt_on_missing=False):
 512         """Create rig instances and connect them into a tree."""
 513
 514         assert(self.context.active_object == self.obj)
 515         assert(self.obj.mode == 'OBJECT')
 516
 517         self.stage = 'instantiate'
 518
 519         # Compute the list of bones
 520         bone_list = list_bone_names_depth_first_sorted(self.obj)
 521
 522         self.org_rename_table = {n: n for n in bone_list}
 523
 524         # Construct the rig instances
 525         for name in bone_list:
 526             self.__create_rigs(self.org_rename_table[name], halt_on_missing)
 527
 528         # Connect rigs and bones into a tree
 529         handled = {}
 530
 531         for bone in self.obj.data.bones:
 532             if bone.parent is None:
 533                 self.__build_rig_tree_rec(bone, None, handled)