diff options
Diffstat (limited to 'rigify/utils/rig.py')
-rw-r--r-- | rigify/utils/rig.py | 56 |
1 files changed, 53 insertions, 3 deletions
diff --git a/rigify/utils/rig.py b/rigify/utils/rig.py index bcb3ff74..fa55c1aa 100644 --- a/rigify/utils/rig.py +++ b/rigify/utils/rig.py @@ -22,6 +22,9 @@ import bpy import importlib import importlib.util import os +import re + +from itertools import count from bpy.types import bpy_struct, bpy_prop_array, Constraint @@ -173,16 +176,54 @@ def _generate_properties(lines, prefix, obj, base_class, *, defaults={}, objects lines.append('%s.%s = %r' % (prefix, prop.identifier, cur_value)) -def write_metarig(obj, layers=False, func_name="create", groups=False): +def write_metarig_widgets(obj): + from .widgets import write_widget + + widget_set = set() + + for pbone in obj.pose.bones: + if pbone.custom_shape: + widget_set.add(pbone.custom_shape) + + id_set = set() + widget_map = {} + code = [] + + for widget_obj in widget_set: + ident = re.sub("[^0-9a-zA-Z_]+", "_", widget_obj.name) + + if ident in id_set: + for i in count(1): + if ident+'_'+str(i) not in id_set: + break + + id_set.add(ident) + widget_map[widget_obj] = ident + + code.append(write_widget(widget_obj, name=ident, use_size=False)) + + return widget_map, code + + +def write_metarig(obj, layers=False, func_name="create", groups=False, widgets=False): """ Write a metarig as a python script, this rig is to have all info needed for generating the real rig with rigify. """ code = [] - code.append("import bpy\n\n") - code.append("from mathutils import Color\n\n") + code.append("import bpy\n") + code.append("from mathutils import Color\n") + # Widget object creation functions if requested + if widgets: + widget_map, widget_code = write_metarig_widgets(obj) + + if widget_map: + code.append("from rigify.utils.widgets import widget_generator\n\n") + code += widget_code + + # Start of the metarig function code.append("def %s(obj):" % func_name) code.append(" # generated by rigify.utils.write_metarig") bpy.ops.object.mode_set(mode='EDIT') @@ -247,6 +288,9 @@ def write_metarig(obj, layers=False, func_name="create", groups=False): code.append("") code.append(" bpy.ops.object.mode_set(mode='OBJECT')") + if widgets and widget_map: + code.append(" widget_map = {}") + # Rig type and other pose properties for bone_name in bones: pbone = obj.pose.bones[bone_name] @@ -294,6 +338,12 @@ def write_metarig(obj, layers=False, func_name="create", groups=False): }, objects={obj: 'obj'}, ) + # Custom widgets + if widgets and pbone.custom_shape: + widget_id = widget_map[pbone.custom_shape] + code.append(" if %r not in widget_map:" % (widget_id)) + code.append(" widget_map[%r] = create_%s_widget(obj, pbone.name, widget_name=%r, widget_force_new=True)" % (widget_id, widget_id, pbone.custom_shape.name)) + code.append(" pbone.custom_shape = widget_map[%r]" % (widget_id)) code.append("\n bpy.ops.object.mode_set(mode='EDIT')") code.append(" for bone in arm.edit_bones:") |