Skip to content

Commit cb41b0f

Browse files
committed
Rewrite generate.py into a class and introduce a base rig class.
The main goals are to provide an official way for rigs to interact in a structured way, and to remove mode switching within rigs. This involves introducing a base class for rigs that holds rig-to-rig and rig-to-bone references, converting the main generator into a class and passing it to rigs, and splitting the single generate method into multiple passes. For backward compatibility, old rigs are automatically handled via a wrapper that translates between old and new API. In addition, a way to create objects that receive the generate callbacks that aren't rigs is introduced via the GeneratorPlugin class. The UI script generation code is converted into a plugin.
1 parent 61f0cd1 commit cb41b0f

File tree

16 files changed

+1420
-588
lines changed

16 files changed

+1420
-588
lines changed

__init__.py

Lines changed: 9 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -32,14 +32,17 @@
3232

3333
if "bpy" in locals():
3434
import importlib
35+
importlib.reload(utils)
36+
importlib.reload(base_rig)
37+
importlib.reload(base_generate)
38+
importlib.reload(rig_ui_template)
39+
importlib.reload(rig_lists)
3540
importlib.reload(generate)
3641
importlib.reload(ui)
37-
importlib.reload(utils)
3842
importlib.reload(metarig_menu)
39-
importlib.reload(rig_lists)
4043
importlib.reload(feature_sets)
4144
else:
42-
from . import (utils, rig_lists, generate, ui, metarig_menu, feature_sets)
45+
from . import (utils, base_rig, base_generate, rig_ui_template, rig_lists, generate, ui, metarig_menu, feature_sets)
4346

4447
import bpy
4548
import sys
@@ -449,7 +452,9 @@ def update_mode(self, context):
449452
pass
450453
else:
451454
for rig in rig_lists.rigs:
452-
r = rig_lists.rigs[rig]['module']
455+
rig_module = rig_lists.rigs[rig]['module']
456+
rig_class = rig_module.Rig
457+
r = rig_class if issubclass(rig_class, base_rig.BaseRig) else rig_module
453458
try:
454459
r.add_parameters(RigifyParameterValidator(RigifyParameters, rig, RIGIFY_PARAMETER_TABLE))
455460
except AttributeError:

base_generate.py

Lines changed: 353 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,353 @@
1+
#====================== BEGIN GPL LICENSE BLOCK ======================
2+
#
3+
# This program is free software; you can redistribute it and/or
4+
# modify it under the terms of the GNU General Public License
5+
# as published by the Free Software Foundation; either version 2
6+
# of the License, or (at your option) any later version.
7+
#
8+
# This program is distributed in the hope that it will be useful,
9+
# but WITHOUT ANY WARRANTY; without even the implied warranty of
10+
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
11+
# GNU General Public License for more details.
12+
#
13+
# You should have received a copy of the GNU General Public License
14+
# along with this program; if not, write to the Free Software Foundation,
15+
# Inc., 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301, USA.
16+
#
17+
#======================= END GPL LICENSE BLOCK ========================
18+
19+
# <pep8 compliant>
20+
21+
import bpy
22+
import sys
23+
import traceback
24+
25+
from .utils.errors import MetarigError
26+
from .utils.naming import random_id
27+
from .utils.metaclass import SingletonPluginMetaclass
28+
29+
from . import base_rig
30+
31+
32+
#=============================================
33+
# Generator Plugin
34+
#=============================================
35+
36+
37+
class GeneratorPlugin(base_rig.GenerateCallbackMixin, metaclass=SingletonPluginMetaclass):
38+
"""
39+
Base class for generator plugins.
40+
41+
Generator plugins are per-Generator singleton utility
42+
classes that receive the same stage callbacks as rigs.
43+
44+
Useful for building entities shared by multiple rigs
45+
(e.g. the python script), or for making fire-and-forget
46+
utilities that actually require multiple stages to
47+
complete.
48+
49+
This will create only one instance per set of args:
50+
51+
instance = PluginClass(generator, ...init args)
52+
"""
53+
54+
priority = 0
55+
56+
def __init__(self, generator):
57+
self.generator = generator
58+
self.obj = generator.obj
59+
60+
def register_new_bone(self, new_name, old_name=None):
61+
self.generator.bone_owners[new_name] = None
62+
63+
64+
#=============================================
65+
# Legacy Rig Wrapper
66+
#=============================================
67+
68+
69+
class LegacyRig(base_rig.BaseRig):
70+
"""Wrapper around legacy style rigs without a common base class"""
71+
72+
def __init__(self, generator, bone, wrapped_class):
73+
self.wrapped_rig = None
74+
self.wrapped_class = wrapped_class
75+
76+
super(LegacyRig,self).__init__(generator, bone)
77+
78+
def find_org_bones(self, bone):
79+
if not self.wrapped_rig:
80+
self.wrapped_rig = self.wrapped_class(self.obj, self.base_bone, self.params)
81+
82+
# Try to extract the main list of bones - old rigs often have it.
83+
# This is not actually strictly necessary, so failing is OK.
84+
if hasattr(self.wrapped_rig, 'org_bones'):
85+
bones = self.wrapped_rig.org_bones
86+
if isinstance(bones, list):
87+
return bones
88+
89+
return [bone.name]
90+
91+
def generate_bones(self):
92+
# Old rigs only have one generate method, so call it from
93+
# generate_bones, which is the only stage allowed to add bones.
94+
scripts = self.wrapped_rig.generate()
95+
96+
# Switch back to EDIT mode if the rig changed it
97+
if self.obj.mode != 'EDIT':
98+
bpy.ops.object.mode_set(mode='EDIT')
99+
100+
if isinstance(scripts, dict):
101+
if 'script' in scripts:
102+
self.script.add_panel_code(scripts['script'])
103+
if 'imports' in scripts:
104+
self.script.add_imports(scripts['imports'])
105+
if 'utilities' in scripts:
106+
self.script.add_utilities(scripts['utilities'])
107+
if 'register' in scripts:
108+
self.script.register_classes(scripts['register'])
109+
if 'register_drivers' in scripts:
110+
self.script.register_driver_functions(scripts['register_drivers'])
111+
if 'register_props' in scripts:
112+
for prop, val in scripts['register_props']:
113+
self.script.register_property(prop, val)
114+
if 'noparent_bones' in scripts:
115+
for bone in scripts['noparent_bones']:
116+
self.generator.disable_auto_parent(bone)
117+
elif scripts is not None:
118+
self.script.add_panel_code([scripts[0]])
119+
120+
def finalize(self):
121+
if hasattr(self.wrapped_rig, 'glue'):
122+
self.wrapped_rig.glue()
123+
124+
125+
#=============================================
126+
# Base Generate Engine
127+
#=============================================
128+
129+
130+
class BaseGenerator(object):
131+
"""Base class for the main generator object. Contains rig and plugin management code."""
132+
133+
def __init__(self, context, metarig):
134+
self.context = context
135+
self.scene = context.scene
136+
self.metarig = metarig
137+
self.obj = None
138+
139+
# List of all rig instances
140+
self.rig_list = []
141+
# List of rigs that don't have a parent
142+
self.root_rigs = []
143+
# Map from bone names to their rigs
144+
self.bone_owners = {}
145+
146+
# Set of plugins
147+
self.plugin_list = []
148+
self.plugin_map = {}
149+
150+
# Current execution stage so plugins could check they are used correctly
151+
self.stage = None
152+
153+
# Set of bones that should be left without parent
154+
self.noparent_bones = set()
155+
156+
# Random string with time appended so that
157+
# different rigs don't collide id's
158+
self.rig_id = random_id(16)
159+
160+
161+
def disable_auto_parent(self, bone):
162+
"""Prevent automatically parenting the bone to root if parentless."""
163+
self.noparent_bones.add(bone)
164+
165+
166+
def __run_object_stage(self, method_name):
167+
assert(self.context.active_object == self.obj)
168+
assert(self.obj.mode == 'OBJECT')
169+
num_bones = len(self.obj.data.bones)
170+
171+
self.stage = method_name
172+
173+
for rig in [*self.rig_list, *self.plugin_list]:
174+
getattr(rig, method_name)()
175+
176+
assert(self.context.active_object == self.obj)
177+
assert(self.obj.mode == 'OBJECT')
178+
assert(num_bones == len(self.obj.data.bones))
179+
180+
181+
def __run_edit_stage(self, method_name):
182+
assert(self.context.active_object == self.obj)
183+
assert(self.obj.mode == 'EDIT')
184+
num_bones = len(self.obj.data.edit_bones)
185+
186+
self.stage = method_name
187+
188+
for rig in [*self.rig_list, *self.plugin_list]:
189+
getattr(rig, method_name)()
190+
191+
assert(self.context.active_object == self.obj)
192+
assert(self.obj.mode == 'EDIT')
193+
assert(num_bones == len(self.obj.data.edit_bones))
194+
195+
196+
def invoke_initialize(self):
197+
self.__run_object_stage('initialize')
198+
199+
200+
def invoke_prepare_bones(self):
201+
self.__run_edit_stage('prepare_bones')
202+
203+
204+
def __auto_register_bones(self, bones, rig):
205+
"""Find bones just added and not registered by this rig."""
206+
for bone in bones:
207+
name = bone.name
208+
if name not in self.bone_owners:
209+
self.bone_owners[name] = rig
210+
if rig:
211+
rig.rigify_new_bones[name] = None
212+
213+
if not isinstance(rig, LegacyRig):
214+
print("WARNING: rig %s didn't register bone %s\n" % (rig, name))
215+
216+
217+
def invoke_generate_bones(self):
218+
assert(self.context.active_object == self.obj)
219+
assert(self.obj.mode == 'EDIT')
220+
221+
self.stage = 'generate_bones'
222+
223+
for rig in self.rig_list:
224+
rig.generate_bones()
225+
226+
assert(self.context.active_object == self.obj)
227+
assert(self.obj.mode == 'EDIT')
228+
229+
self.__auto_register_bones(self.obj.data.edit_bones, rig)
230+
231+
for plugin in self.plugin_list:
232+
plugin.generate_bones()
233+
234+
assert(self.context.active_object == self.obj)
235+
assert(self.obj.mode == 'EDIT')
236+
237+
self.__auto_register_bones(self.obj.data.edit_bones, None)
238+
239+
240+
def invoke_parent_bones(self):
241+
self.__run_edit_stage('parent_bones')
242+
243+
244+
def invoke_configure_bones(self):
245+
self.__run_object_stage('configure_bones')
246+
247+
248+
def invoke_rig_bones(self):
249+
self.__run_object_stage('rig_bones')
250+
251+
252+
def invoke_generate_widgets(self):
253+
self.__run_object_stage('generate_widgets')
254+
255+
256+
def invoke_finalize(self):
257+
self.__run_object_stage('finalize')
258+
259+
260+
def instantiate_rig(self, rig_class, bone):
261+
if issubclass(rig_class, base_rig.BaseRig):
262+
return rig_class(self, bone)
263+
else:
264+
return LegacyRig(self, bone, rig_class)
265+
266+
267+
def __create_rigs_rec(self, bone, halt_on_missing):
268+
"""Recursively walk bones and create rig instances."""
269+
270+
bone_name = bone.name
271+
child_list = [bone.name for bone in bone.children]
272+
273+
pose_bone = self.obj.pose.bones[bone_name]
274+
275+
rig_type = pose_bone.rigify_type
276+
rig_type = rig_type.replace(" ", "")
277+
278+
if rig_type != "":
279+
try:
280+
rig_class = self.find_rig_class(rig_type)
281+
rig = self.instantiate_rig(rig_class, pose_bone)
282+
283+
assert(self.context.active_object == self.obj)
284+
assert(self.obj.mode == 'OBJECT')
285+
286+
self.rig_list.append(rig)
287+
288+
for org_name in rig.rigify_org_bones:
289+
if org_name in self.bone_owners:
290+
print("CONFLICT: bone %s already claimed by rig %s\n" % (org_name, self.bone_owners[org_name]))
291+
292+
self.bone_owners[org_name] = rig
293+
294+
except ImportError:
295+
message = "Rig Type Missing: python module for type '%s' not found (bone: %s)" % (rig_type, bone_name)
296+
if halt_on_missing:
297+
raise MetarigError(message)
298+
else:
299+
print(message)
300+
print('print_exc():')
301+
traceback.print_exc(file=sys.stdout)
302+
303+
child_list.sort()
304+
305+
for child in child_list:
306+
cbone = self.obj.data.bones[child]
307+
self.__create_rigs_rec(cbone, halt_on_missing)
308+
309+
310+
def __build_rig_tree_rec(self, bone, current_rig):
311+
"""Recursively walk bones and connect rig instances into a tree."""
312+
313+
rig = self.bone_owners.get(bone.name)
314+
315+
if rig:
316+
rig.rigify_parent = current_rig
317+
318+
if current_rig:
319+
current_rig.rigify_children.append(rig)
320+
else:
321+
self.root_rigs.append(rig)
322+
323+
current_rig = rig
324+
else:
325+
if current_rig:
326+
current_rig.rigify_child_bones.add(bone.name)
327+
328+
self.bone_owners[bone.name] = current_rig
329+
330+
for child in bone.children:
331+
self.__build_rig_tree_rec(child, current_rig)
332+
333+
334+
def instantiate_rig_tree(self, halt_on_missing=False):
335+
"""Create rig instances and connect them into a tree."""
336+
337+
assert(self.context.active_object == self.obj)
338+
assert(self.obj.mode == 'OBJECT')
339+
340+
bone_names = [bone.name for bone in self.obj.data.bones]
341+
bone_names.sort()
342+
343+
# Construct the rig instances
344+
for name in bone_names:
345+
bone = self.obj.data.bones[name]
346+
if bone.parent is None:
347+
self.__create_rigs_rec(bone, halt_on_missing)
348+
349+
# Connect rigs and bones into a tree
350+
for bone in self.obj.data.bones:
351+
if bone.parent is None:
352+
self.__build_rig_tree_rec(bone, None)
353+

0 commit comments

Comments
 (0)