Rigify 0.5: basic and animals metarigs. metarig custom multi-menu. minor wgt reposition and shapes

1 files changed, 80 insertions, 25 deletions

diff --git a/rigify/metarig_menu.py b/rigify/metarig_menu.py
index 4bdf2701..0c917c40 100644
--- a/rigify/metarig_menu.py
+++ b/rigify/metarig_menu.py
@@ -26,11 +26,23 @@ import bpy
 from . import utils
 
 
-def get_metarig_list(path):
+class ArmatureSubMenu(bpy.types.Menu):
+    # bl_idname = 'ARMATURE_MT_armature_class'
+
+    def draw(self, context):
+        layout = self.layout
+        layout.label(self.bl_label)
+        for op, name in self.operators:
+            text = capwords(name.replace("_", " ")) + " (Meta-Rig)"
+            layout.operator(op, icon='OUTLINER_OB_ARMATURE', text=text)
+
+
+def get_metarig_list(path, depth=0):
     """ Searches for metarig modules, and returns a list of the
         imported modules.
     """
     metarigs = []
+    metarigs_dict = dict()
     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)
@@ -39,8 +51,12 @@ def get_metarig_list(path):
 
     for f in files:
         # Is it a directory?
-        if os.path.isdir(os.path.join(SEARCH_DIR_ABS, f)):
-            continue
+        complete_path = os.path.join(SEARCH_DIR_ABS, f)
+        if os.path.isdir(complete_path) and depth == 0:
+            if f[0] != '_':
+                metarigs_dict[f] = get_metarig_list(f, depth=1)
+            else:
+                continue
         elif not f.endswith(".py"):
             continue
         elif f == "__init__.py":
@@ -48,10 +64,18 @@ def get_metarig_list(path):
         else:
             module_name = f[:-3]
             try:
-                metarigs += [utils.get_metarig_module(module_name)]
+                if depth == 1:
+                    metarigs += [utils.get_metarig_module(module_name, utils.METARIG_DIR + '.' + path)]
+                else:
+                    metarigs += [utils.get_metarig_module(module_name, utils.METARIG_DIR)]
             except (ImportError):
                 pass
-    return metarigs
+
+    if depth == 1:
+        return metarigs
+
+    metarigs_dict[utils.METARIG_DIR] = metarigs
+    return metarigs_dict
 
 
 def make_metarig_add_execute(m):
@@ -85,41 +109,72 @@ def make_metarig_menu_func(bl_idname, text):
     return metarig_menu
 
 
+def make_submenu_func(bl_idname, text):
+    def metarig_menu(self, context):
+        self.layout.menu(bl_idname, icon='OUTLINER_OB_ARMATURE', text=text)
+    return metarig_menu
+
+
 # Get the metarig modules
-metarigs = get_metarig_list("")
+metarigs_dict = get_metarig_list("")
+print(metarigs_dict)
+armature_submenus = []
 
 # Create metarig add Operators
-metarig_ops = []
-for m in metarigs:
-    name = m.__name__.rsplit('.', 1)[1]
+metarig_ops = {}
+for metarig_class in metarigs_dict:
+    metarig_ops[metarig_class] = []
+    for m in metarigs_dict[metarig_class]:
+        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)
+        # 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))
+        metarig_ops[metarig_class].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)]
-
+for metarig_class in metarigs_dict:
+    # Create menu functions
+    if metarig_class != utils.METARIG_DIR:
+        armature_submenus.append(type('Class_' + metarig_class + '_submenu', (ArmatureSubMenu,), {}))
+        armature_submenus[-1].bl_label = metarig_class + ' (submenu)'
+        armature_submenus[-1].bl_idname = 'ARMATURE_MT_%s_class' % metarig_class
+        armature_submenus[-1].operators = []
+        menu_funcs += [make_submenu_func(armature_submenus[-1].bl_idname, metarig_class)]
+
+    for mop, name in metarig_ops[metarig_class]:
+        print(metarig_class)
+        print(metarig_ops[metarig_class])
+        if metarig_class != utils.METARIG_DIR:
+            arm_sub = next((e for e in armature_submenus if e.bl_label == metarig_class + ' (submenu)'), '')
+            arm_sub.operators.append((mop.bl_idname, name,))
+        else:
+            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 cl in metarig_ops:
+        for mop, name in metarig_ops[cl]:
+            bpy.utils.register_class(mop)
+
+    for arm_sub in armature_submenus:
+        bpy.utils.register_class(arm_sub)
 
     for mf in menu_funcs:
         bpy.types.INFO_MT_armature_add.append(mf)
 
 
 def unregister():
-    for mop, name in metarig_ops:
-        bpy.utils.unregister_class(mop)
+    for cl in metarig_ops:
+        for mop, name in metarig_ops[cl]:
+            bpy.utils.unregister_class(mop)
+
+    for arm_sub in armature_submenus:
+        bpy.utils.unregister_class(arm_sub)
 
     for mf in menu_funcs:
         bpy.types.INFO_MT_armature_add.remove(mf)