Rigify: the pre-built metarigs that appear in the add armature menu

are now dynamically pulled from python files in the metarigs
directory.  This makes adding new metarigs pretty painless.

2 files changed, 99 insertions, 18 deletions

diff --git a/rigify/metarig_menu.py b/rigify/metarig_menu.py
index 6133f37e..d635d50a 100644
--- a/rigify/metarig_menu.py
+++ b/rigify/metarig_menu.py
@@ -18,40 +18,108 @@
 
 # <pep8 compliant>
 
+import os
+from string import capwords
+
 import bpy
-from rigify.metarigs import human
+from . import utils
+
+
+def get_metarig_list(path):
+    """ Searches for metarig modules, and returns a list of the
+        imported modules.
+    """
+    metarigs = []
+    MODULE_DIR = os.path.dirname(__file__)
+    METARIG_DIR_ABS = os.path.join(MODULE_DIR, utils.METARIG_DIR)
+    SEARCH_DIR_ABS = os.path.join(METARIG_DIR_ABS, path)
+    files = os.listdir(SEARCH_DIR_ABS)
+    files.sort()
 
+    for f in files:
+        # Is it a directory?
+        if os.path.isdir(os.path.join(SEARCH_DIR_ABS, f)):
+            continue
+        elif not f.endswith(".py"):
+            continue
+        elif f == "__init__.py":
+            continue
+        else:
+            module_name = f[:-3]
+            try:
+                metarigs += [utils.get_metarig_module(module_name)]
+            except (ImportError):
+                pass
+    return metarigs
 
-class AddHuman(bpy.types.Operator):
-    """Add an advanced human metarig base"""
-    bl_idname = "object.armature_human_advanced_add"
-    bl_label = "Add Humanoid (advanced metarig)"
-    bl_options = {'REGISTER', 'UNDO'}
 
+def make_metarig_add_execute(m):
+    """ Create an execute method for a metarig creation operator.
+    """
     def execute(self, context):
+        # Add armature object
         bpy.ops.object.armature_add()
         obj = context.active_object
-        mode_orig = obj.mode
-        bpy.ops.object.mode_set(mode='EDIT')  # grr, remove bone
+        obj.name = "metarig"
+        
+        # Remove default bone
+        bpy.ops.object.mode_set(mode='EDIT')
         bones = context.active_object.data.edit_bones
         bones.remove(bones[0])
-        human.create(obj)
-        bpy.ops.object.mode_set(mode=mode_orig)
+        
+        # Create metarig
+        m.create(obj)
+        
+        bpy.ops.object.mode_set(mode='OBJECT')
         return {'FINISHED'}
+    return execute
 
 
-# Add to a menu
-menu_func = (lambda self, context: self.layout.operator(AddHuman.bl_idname,
-                    icon='OUTLINER_OB_ARMATURE', text="Human (Meta-Rig)"))
+def make_metarig_menu_func(bl_idname, text):
+    """ For some reason lambda's don't work for adding multiple menu
+        items, so we use this instead to generate the functions.
+    """
+    def metarig_menu(self, context):
+        self.layout.operator(bl_idname, icon='OUTLINER_OB_ARMATURE', text=text)
+    return metarig_menu
 
 
-def register():
-    bpy.utils.register_class(AddHuman)
 
-    bpy.types.INFO_MT_armature_add.append(menu_func)
+# Get the metarig modules
+metarigs = get_metarig_list("")
+
+# Create metarig add Operators
+metarig_ops = []
+for m in metarigs:
+    name = m.__name__.rsplit('.', 1)[1]
+    
+    # Dynamically construct an Operator
+    T = type("Add_" + name + "_Metarig", (bpy.types.Operator,), {})
+    T.bl_idname = "object.armature_" + name + "_metarig_add"
+    T.bl_label = "Add " + name.replace("_", " ").capitalize() + " (metarig)"
+    T.bl_options = {'REGISTER', 'UNDO'}
+    T.execute = make_metarig_add_execute(m)
+        
+    metarig_ops.append((T, name))
+
+# Create menu functions
+menu_funcs = []
+for mop, name in metarig_ops:
+    text = capwords(name.replace("_", " ")) + " (Meta-Rig)"
+    
+    menu_funcs += [make_metarig_menu_func(mop.bl_idname, text)]
+
+
+def register():
+    for mop, name in metarig_ops:
+        bpy.utils.register_class(mop)
 
+    for mf in menu_funcs:
+        bpy.types.INFO_MT_armature_add.append(mf)
 
 def unregister():
-    bpy.utils.unregister_class(AddHuman)
+    for mop in metarig_ops:
+        bpy.utils.unregister_class(mop)
 
-    bpy.types.INFO_MT_armature_add.remove(menu_func)
+    for mf in menu_funcs:
+        bpy.types.INFO_MT_armature_add.remove(mf)
diff --git a/rigify/utils.py b/rigify/utils.py
index 379e1ea1..728fc7f9 100644
--- a/rigify/utils.py
+++ b/rigify/utils.py
@@ -26,6 +26,7 @@ from mathutils import Vector
 from rna_prop_ui import rna_idprop_ui_prop_get
 
 RIG_DIR = "rigs"  # Name of the directory where rig types are kept
+METARIG_DIR = "metarigs"  # Name of the directory where metarigs are kept
 
 ORG_PREFIX = "ORG-"  # Prefix of original bones.
 MCH_PREFIX = "MCH-"  # Prefix of mechanism bones.
@@ -421,6 +422,18 @@ def get_rig_type(rig_type):
     return submod
 
 
+def get_metarig_module(metarig):
+    """ Fetches a rig module by name, and returns it.
+    """
+    #print("%s.%s.%s" % (__package__,METARIG_DIR,metarig))
+    name="%s.%s.%s" % (MODULE_NAME, METARIG_DIR, metarig)
+    submod = __import__(name)
+    for c in (name.split("."))[1:]:
+        submod = getattr(submod, c)
+    imp.reload(submod)
+    return submod
+
+
 def connected_children_names(obj, bone_name):
     """ Returns a list of bone names (in order) of the bones that form a single
         connected chain starting with the given bone as a parent.