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 def get_metarig_list(path):
  30     """ Searches for metarig modules, and returns a list of the
  31         imported modules.
  32     """
  33     metarigs = []
  34     MODULE_DIR = os.path.dirname(__file__)
  35     METARIG_DIR_ABS = os.path.join(MODULE_DIR, utils.METARIG_DIR)
  36     SEARCH_DIR_ABS = os.path.join(METARIG_DIR_ABS, path)
  37     files = os.listdir(SEARCH_DIR_ABS)
  38     files.sort()
  39
  40     for f in files:
  41         # Is it a directory?
  42         if os.path.isdir(os.path.join(SEARCH_DIR_ABS, f)):
  43             continue
  44         elif not f.endswith(".py"):
  45             continue
  46         elif f == "__init__.py":
  47             continue
  48         else:
  49             module_name = f[:-3]
  50             try:
  51                 metarigs += [utils.get_metarig_module(module_name)]
  52             except (ImportError):
  53                 pass
  54     return metarigs
  55
  56
  57 def make_metarig_add_execute(m):
  58     """ Create an execute method for a metarig creation operator.
  59     """
  60     def execute(self, context):
  61         # Add armature object
  62         bpy.ops.object.armature_add()
  63         obj = context.active_object
  64         obj.name = "metarig"
  65
  66         # Remove default bone
  67         bpy.ops.object.mode_set(mode='EDIT')
  68         bones = context.active_object.data.edit_bones
  69         bones.remove(bones[0])
  70
  71         # Create metarig
  72         m.create(obj)
  73
  74         bpy.ops.object.mode_set(mode='OBJECT')
  75         return {'FINISHED'}
  76     return execute
  77
  78
  79 def make_metarig_menu_func(bl_idname, text):
  80     """ For some reason lambda's don't work for adding multiple menu
  81         items, so we use this instead to generate the functions.
  82     """
  83     def metarig_menu(self, context):
  84         self.layout.operator(bl_idname, icon='OUTLINER_OB_ARMATURE', text=text)
  85     return metarig_menu
  86
  87
  88 # Get the metarig modules
  89 metarigs = get_metarig_list("")
  90
  91 # Create metarig add Operators
  92 metarig_ops = []
  93 for m in metarigs:
  94     name = m.__name__.rsplit('.', 1)[1]
  95
  96     # Dynamically construct an Operator
  97     T = type("Add_" + name + "_Metarig", (bpy.types.Operator,), {})
  98     T.bl_idname = "object.armature_" + name + "_metarig_add"
  99     T.bl_label = "Add " + name.replace("_", " ").capitalize() + " (metarig)"
 100     T.bl_options = {'REGISTER', 'UNDO'}
 101     T.execute = make_metarig_add_execute(m)
 102
 103     metarig_ops.append((T, name))
 104
 105 # Create menu functions
 106 menu_funcs = []
 107 for mop, name in metarig_ops:
 108     text = capwords(name.replace("_", " ")) + " (Meta-Rig)"
 109     menu_funcs += [make_metarig_menu_func(mop.bl_idname, text)]
 110
 111
 112 def register():
 113     for mop, name in metarig_ops:
 114         bpy.utils.register_class(mop)
 115
 116     for mf in menu_funcs:
 117         bpy.types.INFO_MT_armature_add.append(mf)
 118
 119
 120 def unregister():
 121     for mop, name in metarig_ops:
 122         bpy.utils.unregister_class(mop)
 123
 124     for mf in menu_funcs:
 125         bpy.types.INFO_MT_armature_add.remove(mf)