534 lines
18 KiB
Python
534 lines
18 KiB
Python
# SPDX-License-Identifier: GPL-2.0-or-later
|
|
|
|
import bpy
|
|
import sys
|
|
import traceback
|
|
import collections
|
|
|
|
from typing import Optional, TYPE_CHECKING, Collection, List
|
|
from bpy.types import PoseBone, Bone
|
|
|
|
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, get_rigify_params
|
|
from .utils.misc import clone_parameters, assign_parameters, ArmatureObject
|
|
|
|
from . import base_rig
|
|
|
|
from itertools import count
|
|
|
|
if TYPE_CHECKING:
|
|
from .rig_ui_template import ScriptGenerator
|
|
|
|
|
|
##############################################
|
|
# 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: 'BaseGenerator'):
|
|
self.generator = generator
|
|
self.obj = generator.obj
|
|
|
|
def register_new_bone(self, new_name: str, old_name: Optional[str] = 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: 'BaseGenerator', pose_bone: PoseBone):
|
|
self.generator = generator
|
|
|
|
self.obj = generator.obj
|
|
self.base_bone = pose_bone.name
|
|
self.params = get_rigify_params(pose_bone)
|
|
self.params_copy = clone_parameters(self.params)
|
|
|
|
def substitute(self):
|
|
# return [rig1, rig2...]
|
|
raise NotImplementedError
|
|
|
|
# Utility methods
|
|
def register_new_bone(self, new_name: str, old_name: Optional[str] = None):
|
|
pass
|
|
|
|
def get_params(self, bone_name: str):
|
|
return get_rigify_params(self.obj.pose.bones[bone_name])
|
|
|
|
def assign_params(self, bone_name: str, param_dict=None, **params):
|
|
assign_parameters(self.get_params(bone_name), param_dict, **params)
|
|
|
|
def instantiate_rig(self, rig_class: str | type, bone_name: str):
|
|
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: 'BaseGenerator', pose_bone: PoseBone, wrapped_class: type):
|
|
self.wrapped_rig = None
|
|
self.wrapped_class = wrapped_class
|
|
|
|
super().__init__(generator, pose_bone)
|
|
|
|
def find_org_bones(self, pose_bone: PoseBone):
|
|
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: Optional['BaseGenerator'] = None # static
|
|
|
|
context: bpy.types.Context
|
|
scene: bpy.types.Scene
|
|
view_layer: bpy.types.ViewLayer
|
|
layer_collection: bpy.types.LayerCollection
|
|
collection: bpy.types.Collection
|
|
|
|
metarig: ArmatureObject
|
|
obj: ArmatureObject
|
|
|
|
script: 'ScriptGenerator'
|
|
|
|
rig_list: List[base_rig.BaseRig]
|
|
root_rigs: List[base_rig.BaseRig]
|
|
|
|
bone_owners: dict[str, Optional[base_rig.BaseRig]]
|
|
derived_bones: dict[str, set[str]]
|
|
|
|
stage: Optional[str]
|
|
rig_id: str
|
|
|
|
widget_collection: bpy.types.Collection
|
|
use_mirror_widgets: bool
|
|
old_widget_table: dict[str, bpy.types.Object]
|
|
new_widget_table: dict[str, bpy.types.Object]
|
|
widget_mirror_mesh: dict[str, bpy.types.Mesh]
|
|
|
|
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
|
|
|
|
# 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: str):
|
|
"""Prevent automatically parenting the bone to root if parentless."""
|
|
self.noparent_bones.add(bone_name)
|
|
|
|
def find_derived_bones(self, bone_name: str, *, by_owner=False, recursive=True) -> set[str]:
|
|
"""Find which bones were copied from the specified one."""
|
|
if by_owner:
|
|
owner = self.bone_owners.get(bone_name, None)
|
|
if not owner:
|
|
return set()
|
|
|
|
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: str,
|
|
layers: Collection[bool], priority: float):
|
|
for i, val in enumerate(layers):
|
|
if val:
|
|
self.layer_group_priorities[bone_name][i] = priority
|
|
|
|
def rename_org_bone(self, old_name: str, new_name: str) -> str:
|
|
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: str):
|
|
"""Run a generation stage in Object mode."""
|
|
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: str):
|
|
"""Run a generation stage in Edit mode."""
|
|
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(f"WARNING: rig {self.describe_rig(rig)} "
|
|
f"didn't register bone {name}\n")
|
|
else:
|
|
print(f"WARNING: plugin {plugin} didn't register bone {name}\n")
|
|
|
|
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: type, pose_bone: PoseBone) -> base_rig.BaseRig:
|
|
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 find_rig_class(self, rig_type: str) -> type:
|
|
raise NotImplementedError
|
|
|
|
def instantiate_rig_by_type(self, rig_type: str, pose_bone: PoseBone):
|
|
return self.instantiate_rig(self.find_rig_class(rig_type), pose_bone)
|
|
|
|
# noinspection PyMethodMayBeStatic
|
|
def describe_rig(self, rig: base_rig.BaseRig) -> str:
|
|
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(f"CONFLICT: bone {org_name} is claimed by rigs "
|
|
f"{old_rig} and {new_rig}\n")
|
|
|
|
self.bone_owners[org_name] = rig
|
|
|
|
except ImportError:
|
|
message = f"Rig Type Missing: python module for type '{rig_type}' "\
|
|
f"not found (bone: {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: Bone, current_rig: Optional[base_rig.BaseRig],
|
|
handled: dict[base_rig.BaseRig, str]):
|
|
"""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 {bone.name} owned by rig {rig.base_bone} "
|
|
f"has different parent rig from {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)
|