rigify/base_generate.py

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