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.

4 files changed, 118 insertions, 32 deletions

diff --git a/rigify/generate.py b/rigify/generate.py
index 475c57fe..f1586873 100644
--- a/rigify/generate.py
+++ b/rigify/generate.py
@@ -270,18 +270,25 @@ class Generator(base_generate.BaseGenerator):
         # others make non-deforming.
         for bone in bones:
             name = bone.name
+            layers = None
 
             bone.use_deform = name.startswith(DEF_PREFIX)
 
             # Move all the original bones to their layer.
             if name.startswith(ORG_PREFIX):
-                bone.layers = ORG_LAYER
+                layers = ORG_LAYER
             # Move all the bones with names starting with "MCH-" to their layer.
             elif name.startswith(MCH_PREFIX):
-                bone.layers = MCH_LAYER
+                layers = MCH_LAYER
             # Move all the bones with names starting with "DEF-" to their layer.
             elif name.startswith(DEF_PREFIX):
-                bone.layers = DEF_LAYER
+                layers = DEF_LAYER
+
+            if layers is not None:
+                bone.layers = layers
+
+                # Remove custom shapes from non-control bones
+                bone.custom_shape = None
 
             bone.bbone_x = bone.bbone_z = bone.length * 0.05
 
diff --git a/rigify/ui.py b/rigify/ui.py
index 49c11aaf..7b6539cd 100644
--- a/rigify/ui.py
+++ b/rigify/ui.py
@@ -858,7 +858,7 @@ class EncodeMetarig(bpy.types.Operator):
         else:
             text_block = bpy.data.texts.new(name)
 
-        text = write_metarig(context.active_object, layers=True, func_name="create", groups=True)
+        text = write_metarig(context.active_object, layers=True, func_name="create", groups=True, widgets=True)
         text_block.write(text)
         bpy.ops.object.mode_set(mode='EDIT')
 
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:")
diff --git a/rigify/utils/widgets.py b/rigify/utils/widgets.py
index f198afda..bc9070d4 100644
--- a/rigify/utils/widgets.py
+++ b/rigify/utils/widgets.py
@@ -20,6 +20,7 @@
 
 import bpy
 import math
+import functools
 
 from mathutils import Matrix
 
@@ -94,6 +95,39 @@ def create_widget(rig, bone_name, bone_transform_name=None, *, widget_name=None,
     return obj
 
 
+class GeometryData:
+    def __init__(self):
+        self.verts = []
+        self.edges = []
+        self.faces = []
+
+
+def widget_generator(generate_func):
+    """
+    Decorator that encapsulates a call to create_widget, and only requires
+    the actual function to fill the provided vertex and edge lists.
+
+    Accepts parameters of create_widget, plus any keyword arguments the
+    wrapped function has.
+    """
+    @functools.wraps(generate_func)
+    def wrapper(rig, bone_name, bone_transform_name=None, widget_name=None, widget_force_new=False, **kwargs):
+        obj = create_widget(rig, bone_name, bone_transform_name, widget_name=widget_name, widget_force_new=widget_force_new)
+        if obj is not None:
+            geom = GeometryData()
+
+            generate_func(geom, **kwargs)
+
+            mesh = obj.data
+            mesh.from_pydata(geom.verts, geom.edges, geom.faces)
+            mesh.update()
+            return obj
+        else:
+            return None
+
+    return wrapper
+
+
 def create_circle_polygon(number_verts, axis, radius=1.0, head_tail=0.0):
     """ Creates a basic circle around of an axis selected.
         number_verts: number of vertices of the polygon
@@ -174,44 +208,39 @@ def adjust_widget_transform_mesh(obj, matrix, local=None):
         obj.data.transform(matrix)
 
 
-def write_widget(obj):
+def write_widget(obj, name='thing', use_size=True):
     """ Write a mesh object as a python script for widget use.
     """
     script = ""
-    script += "def create_thing_widget(rig, bone_name, size=1.0, bone_transform_name=None):\n"
-    script += "    obj = create_widget(rig, bone_name, bone_transform_name)\n"
-    script += "    if obj is not None:\n"
+    script += "@widget_generator\n"
+    script += "def create_"+name+"_widget(geom";
+    if use_size:
+        script += ", *, size=1.0"
+    script += "):\n"
 
     # Vertices
-    script += "        verts = ["
+    szs = "*size" if use_size else ""
+    width = 2 if use_size else 3
+
+    script += "    geom.verts = ["
     for i, v in enumerate(obj.data.vertices):
-        script += "({:g}*size, {:g}*size, {:g}*size),".format(v.co[0], v.co[1], v.co[2])
-        script += "\n                 " if i % 2 == 1 else " "
+        script += "({:g}{}, {:g}{}, {:g}{}),".format(v.co[0], szs, v.co[1], szs, v.co[2], szs)
+        script += "\n                  " if i % width == (width - 1) else " "
     script += "]\n"
 
     # Edges
-    script += "        edges = ["
+    script += "    geom.edges = ["
     for i, e in enumerate(obj.data.edges):
         script += "(" + str(e.vertices[0]) + ", " + str(e.vertices[1]) + "),"
-        script += "\n                 " if i % 10 == 9 else " "
+        script += "\n                  " if i % 10 == 9 else " "
     script += "]\n"
 
     # Faces
-    script += "        faces = ["
-    for i, f in enumerate(obj.data.polygons):
-        script += "("
-        for v in f.vertices:
-            script += str(v) + ", "
-        script += "),"
-        script += "\n                 " if i % 10 == 9 else " "
-    script += "]\n"
-
-    # Build mesh
-    script += "\n        mesh = obj.data\n"
-    script += "        mesh.from_pydata(verts, edges, faces)\n"
-    script += "        mesh.update()\n"
-    script += "        return obj\n"
-    script += "    else:\n"
-    script += "        return None\n"
+    if obj.data.polygons:
+        script += "    geom.faces = ["
+        for i, f in enumerate(obj.data.polygons):
+            script += "(" + ", ".join(str(v) for v in f.vertices) + "),"
+            script += "\n                  " if i % 10 == 9 else " "
+        script += "]\n"
 
     return script