diff --git a/rigify/generate.py b/rigify/generate.py
index 475c57feb2a1828c1b1b3a95ef63140c29bd2f78..f1586873698f797f244915136b0301583db0a289 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 49c11aafae8220dd126c2ebc23ef375736b07e8b..7b6539cdc7a7c7509a38b60627086831ee257469 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 bcb3ff74ed293b16587fb3e84bedb4b38ad52f08..fa55c1aa3a6ca3eb75963314e8e42ffbe380930a 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 f198afdaca03c3fb5a3dd45bbde1a46ba641052e..bc9070d4d03eb01f6b36171891cf352b40830d9c 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