Rigify: include widgets in generated metarig code.

Since rigs like super_copy already support using widgets assigned
directly to metarig bones, implement adding them with the metarig.

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:")