Rigify: improved rigify.utils.write_metarig()

It is now able to write out all relevant layer data.
Also, the script generated can now be directly run, for testing.

1 files changed, 32 insertions, 1 deletions

diff --git a/rigify/utils.py b/rigify/utils.py
index 728fc7f9..3e23e3a4 100644
--- a/rigify/utils.py
+++ b/rigify/utils.py
@@ -494,12 +494,14 @@ def get_layers(layers):
             return [x in layers for x in range(0, 32)]
 
 
-def write_metarig(obj, layers=False, func_name="create_sample"):
+def write_metarig(obj, layers=False, func_name="create"):
     """
     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")
 
     code.append("def %s(obj):" % func_name)
     code.append("    # generated by rigify.utils.write_metarig")
@@ -508,6 +510,18 @@ def write_metarig(obj, layers=False, func_name="create_sample"):
     code.append("    arm = obj.data")
 
     arm = obj.data
+
+    # Rigify layer layout info
+    if layers and len(arm.rigify_layers) > 0:
+        code.append("\n    for i in range(" + str(len(arm.rigify_layers)) + "):")
+        code.append("        arm.rigify_layers.add()\n")
+
+        for i in range(len(arm.rigify_layers)):
+            name = arm.rigify_layers[i].name
+            row = arm.rigify_layers[i].row
+            code.append('    arm.rigify_layers[' + str(i) + '].name = "' + name + '"')
+            code.append('    arm.rigify_layers[' + str(i) + '].row = ' + str(row))
+
     # write parents first
     bones = [(len(bone.parent_recursive), bone.name) for bone in arm.edit_bones]
     bones.sort(key=lambda item: item[0])
@@ -569,6 +583,23 @@ def write_metarig(obj, layers=False, func_name="create_sample"):
     code.append("        bone.select_tail = True")
     code.append("        arm.edit_bones.active = bone")
 
+    # Set appropriate layers visible
+    if layers:
+        # Find what layers have bones on them
+        active_layers = []
+        for bone_name in bones:
+            bone = obj.data.bones[bone_name]
+            for i in range(len(bone.layers)):
+                if bone.layers[i]:
+                    if i not in active_layers:
+                        active_layers.append(i)
+        active_layers.sort()
+
+        code.append("\n    arm.layers = [(x in " + str(active_layers) + ") for x in range(" + str(len(arm.layers)) + ")]")
+
+    code.append('\nif __name__ == "__main__":')
+    code.append("    " + func_name + "(bpy.context.active_object)")
+
     return "\n".join(code)