# SPDX-License-Identifier: GPL-2.0-or-later import bpy import sys import traceback import collections from .utils.errors import MetarigError, RaiseErrorMixin from .utils.naming import random_id from .utils.metaclass import SingletonPluginMetaclass from .utils.rig import list_bone_names_depth_first_sorted, get_rigify_type from .utils.misc import clone_parameters, assign_parameters from . import base_rig from itertools import count #============================================= # Generator Plugin #============================================= class GeneratorPlugin(base_rig.GenerateCallbackHost, metaclass=SingletonPluginMetaclass): """ Base class for generator plugins. Generator plugins are per-Generator singleton utility classes that receive the same stage callbacks as rigs. Useful for building entities shared by multiple rigs (e.g. the python script), or for making fire-and-forget utilities that actually require multiple stages to complete. This will create only one instance per set of args: instance = PluginClass(generator, ...init args) """ priority = 0 def __init__(self, generator): self.generator = generator self.obj = generator.obj def register_new_bone(self, new_name, old_name=None): self.generator.bone_owners[new_name] = None if old_name: self.generator.derived_bones[old_name].add(new_name) #============================================= # Rig Substitution Mechanism #============================================= class SubstitutionRig(RaiseErrorMixin): """A proxy rig that replaces itself with one or more different rigs.""" def __init__(self, generator, pose_bone): self.generator = generator self.obj = generator.obj self.base_bone = pose_bone.name self.params = pose_bone.rigify_parameters self.params_copy = clone_parameters(self.params) def substitute(self): # return [rig1, rig2...] raise NotImplementedException() # Utility methods def register_new_bone(self, new_name, old_name=None): pass def get_params(self, bone_name): return self.obj.pose.bones[bone_name].rigify_parameters def assign_params(self, bone_name, param_dict=None, **params): assign_parameters(self.get_params(bone_name), param_dict, **params) def instantiate_rig(self, rig_class, bone_name): if isinstance(rig_class, str): rig_class = self.generator.find_rig_class(rig_class) return self.generator.instantiate_rig(rig_class, self.obj.pose.bones[bone_name]) #============================================= # Legacy Rig Wrapper #============================================= class LegacyRig(base_rig.BaseRig): """Wrapper around legacy style rigs without a common base class""" def __init__(self, generator, pose_bone, wrapped_class): self.wrapped_rig = None self.wrapped_class = wrapped_class super().__init__(generator, pose_bone) def find_org_bones(self, pose_bone): bone_name = pose_bone.name if not self.wrapped_rig: self.wrapped_rig = self.wrapped_class(self.obj, self.base_bone, self.params) # Switch back to OBJECT mode if the rig changed it if self.obj.mode != 'OBJECT': bpy.ops.object.mode_set(mode='OBJECT') # Try to extract the main list of bones - old rigs often have it. # This is not actually strictly necessary, so failing is OK. if hasattr(self.wrapped_rig, 'org_bones'): bones = self.wrapped_rig.org_bones if isinstance(bones, list): return bones return [bone_name] def generate_bones(self): # Inject references into the rig if it won't cause conflict if not hasattr(self.wrapped_rig, 'rigify_generator'): self.wrapped_rig.rigify_generator = self.generator if not hasattr(self.wrapped_rig, 'rigify_wrapper'): self.wrapped_rig.rigify_wrapper = self # Old rigs only have one generate method, so call it from # generate_bones, which is the only stage allowed to add bones. scripts = self.wrapped_rig.generate() # Switch back to EDIT mode if the rig changed it if self.obj.mode != 'EDIT': bpy.ops.object.mode_set(mode='EDIT') if isinstance(scripts, dict): if 'script' in scripts: self.script.add_panel_code(scripts['script']) if 'imports' in scripts: self.script.add_imports(scripts['imports']) if 'utilities' in scripts: self.script.add_utilities(scripts['utilities']) if 'register' in scripts: self.script.register_classes(scripts['register']) if 'register_drivers' in scripts: self.script.register_driver_functions(scripts['register_drivers']) if 'register_props' in scripts: for prop, val in scripts['register_props']: self.script.register_property(prop, val) if 'noparent_bones' in scripts: for bone_name in scripts['noparent_bones']: self.generator.disable_auto_parent(bone_name) elif scripts is not None: self.script.add_panel_code([scripts[0]]) def finalize(self): if hasattr(self.wrapped_rig, 'glue'): self.wrapped_rig.glue() # Switch back to OBJECT mode if the rig changed it if self.obj.mode != 'OBJECT': bpy.ops.object.mode_set(mode='OBJECT') #============================================= # Base Generate Engine #============================================= class BaseGenerator: """Base class for the main generator object. Contains rig and plugin management code.""" instance = None def __init__(self, context, metarig): self.context = context self.scene = context.scene self.view_layer = context.view_layer self.layer_collection = context.layer_collection self.collection = self.layer_collection.collection self.metarig = metarig self.obj = None # List of all rig instances self.rig_list = [] # List of rigs that don't have a parent self.root_rigs = [] # Map from bone names to their rigs self.bone_owners = {} self.derived_bones = collections.defaultdict(set) # Set of plugins self.plugin_list = [] self.plugin_map = {} # Current execution stage so plugins could check they are used correctly self.stage = None # Set of bones that should be left without parent self.noparent_bones = set() # Table of layer priorities for defining bone groups self.layer_group_priorities = collections.defaultdict(dict) # Random string with time appended so that # different rigs don't collide id's self.rig_id = random_id(16) # Table of renamed ORG bones self.org_rename_table = dict() def disable_auto_parent(self, bone_name): """Prevent automatically parenting the bone to root if parentless.""" self.noparent_bones.add(bone_name) def find_derived_bones(self, bone_name, *, by_owner=False, recursive=True): """Find which bones were copied from the specified one.""" if by_owner: owner = self.bone_owners.get(bone_name, None) if not owner: return {} table = owner.rigify_derived_bones else: table = self.derived_bones if recursive: result = set() def rec(name): for child in table.get(name, {}): result.add(child) rec(child) rec(bone_name) return result else: return set(table.get(bone_name, {})) def set_layer_group_priority(self, bone_name, layers, priority): for i, val in enumerate(layers): if val: self.layer_group_priorities[bone_name][i] = priority def rename_org_bone(self, old_name, new_name): assert self.stage == 'instantiate' assert old_name == self.org_rename_table.get(old_name, None) assert old_name not in self.bone_owners bone = self.obj.data.bones[old_name] bone.name = new_name new_name = bone.name self.org_rename_table[old_name] = new_name return new_name def __run_object_stage(self, method_name): assert(self.context.active_object == self.obj) assert(self.obj.mode == 'OBJECT') num_bones = len(self.obj.data.bones) self.stage = method_name for rig in self.rig_list: rig.rigify_invoke_stage(method_name) assert(self.context.active_object == self.obj) assert(self.obj.mode == 'OBJECT') assert(num_bones == len(self.obj.data.bones)) # Allow plugins to be added to the end of the list on the fly for i in count(0): if i >= len(self.plugin_list): break self.plugin_list[i].rigify_invoke_stage(method_name) assert(self.context.active_object == self.obj) assert(self.obj.mode == 'OBJECT') assert(num_bones == len(self.obj.data.bones)) def __run_edit_stage(self, method_name): assert(self.context.active_object == self.obj) assert(self.obj.mode == 'EDIT') num_bones = len(self.obj.data.edit_bones) self.stage = method_name for rig in self.rig_list: rig.rigify_invoke_stage(method_name) assert(self.context.active_object == self.obj) assert(self.obj.mode == 'EDIT') assert(num_bones == len(self.obj.data.edit_bones)) # Allow plugins to be added to the end of the list on the fly for i in count(0): if i >= len(self.plugin_list): break self.plugin_list[i].rigify_invoke_stage(method_name) assert(self.context.active_object == self.obj) assert(self.obj.mode == 'EDIT') assert(num_bones == len(self.obj.data.edit_bones)) def invoke_initialize(self): self.__run_object_stage('initialize') def invoke_prepare_bones(self): self.__run_edit_stage('prepare_bones') def __auto_register_bones(self, bones, rig, plugin=None): """Find bones just added and not registered by this rig.""" for bone in bones: name = bone.name if name not in self.bone_owners: self.bone_owners[name] = rig if rig: rig.rigify_new_bones[name] = None if not isinstance(rig, LegacyRig): print("WARNING: rig %s didn't register bone %s\n" % (self.describe_rig(rig), name)) else: print("WARNING: plugin %s didn't register bone %s\n" % (plugin, name)) def invoke_generate_bones(self): assert(self.context.active_object == self.obj) assert(self.obj.mode == 'EDIT') self.stage = 'generate_bones' for rig in self.rig_list: rig.rigify_invoke_stage('generate_bones') assert(self.context.active_object == self.obj) assert(self.obj.mode == 'EDIT') self.__auto_register_bones(self.obj.data.edit_bones, rig) # Allow plugins to be added to the end of the list on the fly for i in count(0): if i >= len(self.plugin_list): break self.plugin_list[i].rigify_invoke_stage('generate_bones') assert(self.context.active_object == self.obj) assert(self.obj.mode == 'EDIT') self.__auto_register_bones(self.obj.data.edit_bones, None, plugin=self.plugin_list[i]) def invoke_parent_bones(self): self.__run_edit_stage('parent_bones') def invoke_configure_bones(self): self.__run_object_stage('configure_bones') def invoke_preapply_bones(self): self.__run_object_stage('preapply_bones') def invoke_apply_bones(self): self.__run_edit_stage('apply_bones') def invoke_rig_bones(self): self.__run_object_stage('rig_bones') def invoke_generate_widgets(self): self.__run_object_stage('generate_widgets') def invoke_finalize(self): self.__run_object_stage('finalize') def instantiate_rig(self, rig_class, pose_bone): assert not issubclass(rig_class, SubstitutionRig) if issubclass(rig_class, base_rig.BaseRig): return rig_class(self, pose_bone) else: return LegacyRig(self, pose_bone, rig_class) def instantiate_rig_by_type(self, rig_type, pose_bone): return self.instantiate_rig(self.find_rig_class(rig_type), pose_bone) def describe_rig(self, rig): base_bone = rig.base_bone if isinstance(rig, LegacyRig): rig = rig.wrapped_rig return "%s (%s)" % (rig.__class__, base_bone) def __create_rigs(self, bone_name, halt_on_missing): """Recursively walk bones and create rig instances.""" pose_bone = self.obj.pose.bones[bone_name] rig_type = get_rigify_type(pose_bone) if rig_type != "": try: rig_class = self.find_rig_class(rig_type) if issubclass(rig_class, SubstitutionRig): rigs = rig_class(self, pose_bone).substitute() else: rigs = [self.instantiate_rig(rig_class, pose_bone)] assert(self.context.active_object == self.obj) assert(self.obj.mode == 'OBJECT') for rig in rigs: self.rig_list.append(rig) for org_name in rig.rigify_org_bones: if org_name in self.bone_owners: old_rig = self.describe_rig(self.bone_owners[org_name]) new_rig = self.describe_rig(rig) print("CONFLICT: bone %s is claimed by rigs %s and %s\n" % (org_name, old_rig, new_rig)) self.bone_owners[org_name] = rig except ImportError: message = "Rig Type Missing: python module for type '%s' not found (bone: %s)" % (rig_type, bone_name) if halt_on_missing: raise MetarigError(message) else: print(message) print('print_exc():') traceback.print_exc(file=sys.stdout) def __build_rig_tree_rec(self, bone, current_rig, handled): """Recursively walk bones and connect rig instances into a tree.""" rig = self.bone_owners.get(bone.name) if rig: if rig is current_rig: pass elif rig not in handled: rig.rigify_parent = current_rig if current_rig: current_rig.rigify_children.append(rig) else: self.root_rigs.append(rig) handled[rig] = bone.name elif rig.rigify_parent is not current_rig: raise MetarigError("CONFLICT: bone %s owned by rig %s has different parent rig from %s\n" % (bone.name, rig.base_bone, handled[rig])) current_rig = rig else: if current_rig: current_rig.rigify_child_bones.add(bone.name) self.bone_owners[bone.name] = current_rig for child in bone.children: self.__build_rig_tree_rec(child, current_rig, handled) def instantiate_rig_tree(self, halt_on_missing=False): """Create rig instances and connect them into a tree.""" assert(self.context.active_object == self.obj) assert(self.obj.mode == 'OBJECT') self.stage = 'instantiate' # Compute the list of bones bone_list = list_bone_names_depth_first_sorted(self.obj) self.org_rename_table = {n: n for n in bone_list} # Construct the rig instances for name in bone_list: self.__create_rigs(self.org_rename_table[name], halt_on_missing) # Connect rigs and bones into a tree handled = {} for bone in self.obj.data.bones: if bone.parent is None: self.__build_rig_tree_rec(bone, None, handled)