rigify/metarig_menu.py

   1 # ##### BEGIN GPL LICENSE BLOCK #####
   2 #
   3 #  This program is free software; you can redistribute it and/or
   4 #  modify it under the terms of the GNU General Public License
   5 #  as published by the Free Software Foundation; either version 2
   6 #  of the License, or (at your option) any later version.
   7 #
   8 #  This program is distributed in the hope that it will be useful,
   9 #  but WITHOUT ANY WARRANTY; without even the implied warranty of
  10 #  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
  11 #  GNU General Public License for more details.
  12 #
  13 #  You should have received a copy of the GNU General Public License
  14 #  along with this program; if not, write to the Free Software Foundation,
  15 #  Inc., 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301, USA.
  16 #
  17 # ##### END GPL LICENSE BLOCK #####
  18
  19 # <pep8 compliant>
  20
  21 import os
  22 from string import capwords
  23
  24 import bpy
  25
  26 from . import utils
  27
  28
  29 class ArmatureSubMenu(bpy.types.Menu):
  30     # bl_idname = 'ARMATURE_MT_armature_class'
  31
  32     def draw(self, context):
  33         layout = self.layout
  34         layout.label(self.bl_label)
  35         for op, name in self.operators:
  36             text = capwords(name.replace("_", " ")) + " (Meta-Rig)"
  37             layout.operator(op, icon='OUTLINER_OB_ARMATURE', text=text)
  38
  39
  40 def get_metarig_list(path, depth=0):
  41     """ Searches for metarig modules, and returns a list of the
  42         imported modules.
  43     """
  44     metarigs = []
  45     metarigs_dict = dict()
  46     MODULE_DIR = os.path.dirname(__file__)
  47     METARIG_DIR_ABS = os.path.join(MODULE_DIR, utils.METARIG_DIR)
  48     SEARCH_DIR_ABS = os.path.join(METARIG_DIR_ABS, path)
  49     files = os.listdir(SEARCH_DIR_ABS)
  50     files.sort()
  51
  52     for f in files:
  53         # Is it a directory?
  54         complete_path = os.path.join(SEARCH_DIR_ABS, f)
  55         if os.path.isdir(complete_path) and depth == 0:
  56             if f[0] != '_':
  57                 metarigs_dict[f] = get_metarig_list(f, depth=1)
  58             else:
  59                 continue
  60         elif not f.endswith(".py"):
  61             continue
  62         elif f == "__init__.py":
  63             continue
  64         else:
  65             module_name = f[:-3]
  66             try:
  67                 if depth == 1:
  68                     metarigs += [utils.get_metarig_module(module_name, utils.METARIG_DIR + '.' + path)]
  69                 else:
  70                     metarigs += [utils.get_metarig_module(module_name, utils.METARIG_DIR)]
  71             except (ImportError):
  72                 pass
  73
  74     if depth == 1:
  75         return metarigs
  76
  77     metarigs_dict[utils.METARIG_DIR] = metarigs
  78     return metarigs_dict
  79
  80
  81 def make_metarig_add_execute(m):
  82     """ Create an execute method for a metarig creation operator.
  83     """
  84     def execute(self, context):
  85         # Add armature object
  86         bpy.ops.object.armature_add()
  87         obj = context.active_object
  88         obj.name = "metarig"
  89         obj.data.name = "metarig"
  90
  91         # Remove default bone
  92         bpy.ops.object.mode_set(mode='EDIT')
  93         bones = context.active_object.data.edit_bones
  94         bones.remove(bones[0])
  95
  96         # Create metarig
  97         m.create(obj)
  98
  99         bpy.ops.object.mode_set(mode='OBJECT')
 100         return {'FINISHED'}
 101     return execute
 102
 103
 104 def make_metarig_menu_func(bl_idname, text):
 105     """ For some reason lambda's don't work for adding multiple menu
 106         items, so we use this instead to generate the functions.
 107     """
 108     def metarig_menu(self, context):
 109         self.layout.operator(bl_idname, icon='OUTLINER_OB_ARMATURE', text=text)
 110     return metarig_menu
 111
 112
 113 def make_submenu_func(bl_idname, text):
 114     def metarig_menu(self, context):
 115         self.layout.menu(bl_idname, icon='OUTLINER_OB_ARMATURE', text=text)
 116     return metarig_menu
 117
 118
 119 # Get the metarig modules
 120 metarigs_dict = get_metarig_list("")
 121 armature_submenus = []
 122
 123 # Create metarig add Operators
 124 metarig_ops = {}
 125 for metarig_class in metarigs_dict:
 126     metarig_ops[metarig_class] = []
 127     for m in metarigs_dict[metarig_class]:
 128         name = m.__name__.rsplit('.', 1)[1]
 129
 130         # Dynamically construct an Operator
 131         T = type("Add_" + name + "_Metarig", (bpy.types.Operator,), {})
 132         T.bl_idname = "object.armature_" + name + "_metarig_add"
 133         T.bl_label = "Add " + name.replace("_", " ").capitalize() + " (metarig)"
 134         T.bl_options = {'REGISTER', 'UNDO'}
 135         T.execute = make_metarig_add_execute(m)
 136
 137         metarig_ops[metarig_class].append((T, name))
 138
 139 menu_funcs = []
 140
 141 for mop, name in metarig_ops[utils.METARIG_DIR]:
 142     text = capwords(name.replace("_", " ")) + " (Meta-Rig)"
 143     menu_funcs += [make_metarig_menu_func(mop.bl_idname, text)]
 144
 145 metarigs_dict.pop(utils.METARIG_DIR)
 146
 147 metarig_classes = list(metarigs_dict.keys())
 148 metarig_classes.sort()
 149 for metarig_class in metarig_classes:
 150     # Create menu functions
 151
 152     armature_submenus.append(type('Class_' + metarig_class + '_submenu', (ArmatureSubMenu,), {}))
 153     armature_submenus[-1].bl_label = metarig_class + ' (submenu)'
 154     armature_submenus[-1].bl_idname = 'ARMATURE_MT_%s_class' % metarig_class
 155     armature_submenus[-1].operators = []
 156     menu_funcs += [make_submenu_func(armature_submenus[-1].bl_idname, metarig_class)]
 157
 158     for mop, name in metarig_ops[metarig_class]:
 159         arm_sub = next((e for e in armature_submenus if e.bl_label == metarig_class + ' (submenu)'), '')
 160         arm_sub.operators.append((mop.bl_idname, name,))
 161
 162 def register():
 163     for cl in metarig_ops:
 164         for mop, name in metarig_ops[cl]:
 165             bpy.utils.register_class(mop)
 166
 167     for arm_sub in armature_submenus:
 168         bpy.utils.register_class(arm_sub)
 169
 170     for mf in menu_funcs:
 171         bpy.types.INFO_MT_armature_add.append(mf)
 172
 173
 174 def unregister():
 175     for cl in metarig_ops:
 176         for mop, name in metarig_ops[cl]:
 177             bpy.utils.unregister_class(mop)
 178
 179     for arm_sub in armature_submenus:
 180         bpy.utils.unregister_class(arm_sub)
 181
 182     for mf in menu_funcs:
 183         bpy.types.INFO_MT_armature_add.remove(mf)