rigify/metarig_menu.py

   1 # SPDX-FileCopyrightText: 2010-2022 Blender Foundation
   2 #
   3 # SPDX-License-Identifier: GPL-2.0-or-later
   4
   5 import os
   6 import traceback
   7
   8 from string import capwords
   9 from collections import defaultdict
  10 from types import ModuleType
  11 from typing import Iterable
  12
  13 import bpy
  14
  15 from .utils.rig import METARIG_DIR, get_resource
  16
  17 from . import feature_set_list
  18
  19
  20 class ArmatureSubMenu(bpy.types.Menu):
  21     # bl_idname = 'ARMATURE_MT_armature_class'
  22
  23     operators: list[tuple[str, str]]
  24
  25     def draw(self, context):
  26         layout = self.layout
  27         layout.label(text=self.bl_label)
  28         for op, name in self.operators:
  29             text = capwords(name.replace("_", " ")) + " (Meta-Rig)"
  30             layout.operator(op, icon='OUTLINER_OB_ARMATURE', text=text)
  31
  32
  33 def get_metarigs(metarig_table: dict[str, ModuleType | dict],
  34                  base_dir: str, base_path: list[str], *,
  35                  path: Iterable[str] = (), nested=False):
  36     """ Searches for metarig modules, and returns a list of the
  37         imported modules.
  38     """
  39
  40     dir_path = os.path.join(base_dir, *path)
  41
  42     try:
  43         files: list[str] = os.listdir(dir_path)
  44     except FileNotFoundError:
  45         files = []
  46
  47     files.sort()
  48
  49     for f in files:
  50         is_dir = os.path.isdir(os.path.join(dir_path, f))  # Whether the file is a directory
  51
  52         # Stop cases
  53         if f[0] in [".", "_"]:
  54             continue
  55         if f.count(".") >= 2 or (is_dir and "." in f):
  56             print("Warning: %r, filename contains a '.', skipping" % os.path.join(*path, f))
  57             continue
  58
  59         if is_dir:  # Check directories
  60             get_metarigs(metarig_table[f], base_dir, base_path, path=[*path, f], nested=True)
  61         elif f.endswith(".py"):
  62             # Check straight-up python files
  63             f = f[:-3]
  64             module = get_resource('.'.join([*base_path, *path, f]))
  65             if nested:
  66                 metarig_table[f] = module
  67             else:
  68                 metarig_table[METARIG_DIR][f] = module
  69
  70
  71 def make_metarig_add_execute(module):
  72     """ Create an execute method for a metarig creation operator.
  73     """
  74     def execute(_self, context):
  75         # Add armature object
  76         bpy.ops.object.armature_add()
  77         obj = context.active_object
  78         obj.name = "metarig"
  79         obj.data.name = "metarig"
  80
  81         # Remove default bone
  82         bpy.ops.object.mode_set(mode='EDIT')
  83         bones = context.active_object.data.edit_bones
  84         bones.remove(bones[0])
  85
  86         # Create metarig
  87         module.create(obj)
  88
  89         bpy.ops.object.mode_set(mode='OBJECT')
  90         return {'FINISHED'}
  91     return execute
  92
  93
  94 def make_metarig_menu_func(bl_idname: str, text: str):
  95     """ For some reason lambdas don't work for adding multiple menu
  96         items, so we use this instead to generate the functions.
  97     """
  98     def metarig_menu(self, _context):
  99         self.layout.operator(bl_idname, icon='OUTLINER_OB_ARMATURE', text=text)
 100     return metarig_menu
 101
 102
 103 def make_submenu_func(bl_idname: str, text: str):
 104     def metarig_menu(self, _context):
 105         self.layout.menu(bl_idname, icon='OUTLINER_OB_ARMATURE', text=text)
 106     return metarig_menu
 107
 108
 109 # Get the metarig modules
 110 def get_internal_metarigs():
 111     base_rigify_dir = os.path.dirname(__file__)
 112     base_rigify_path = __name__.split('.')[:-1]
 113
 114     get_metarigs(metarigs,
 115                  os.path.join(base_rigify_dir, METARIG_DIR),
 116                  [*base_rigify_path, METARIG_DIR])
 117
 118
 119 def infinite_default_dict():
 120     return defaultdict(infinite_default_dict)
 121
 122
 123 metarigs = infinite_default_dict()
 124 metarig_ops = {}
 125 armature_submenus = []
 126 menu_funcs = []
 127
 128
 129 def create_metarig_ops(dic: dict | None = None):
 130     if dic is None:
 131         dic = metarigs
 132
 133     """Create metarig add Operators"""
 134     for metarig_category in dic:
 135         if metarig_category == "external":
 136             create_metarig_ops(dic[metarig_category])
 137             continue
 138         if metarig_category not in metarig_ops:
 139             metarig_ops[metarig_category] = []
 140         for m in dic[metarig_category].values():
 141             name = m.__name__.rsplit('.', 1)[1]
 142
 143             # Dynamically construct an Operator
 144             op_type = type("Add_" + name + "_Metarig", (bpy.types.Operator,), {})
 145             op_type.bl_idname = "object.armature_" + name + "_metarig_add"
 146             op_type.bl_label = "Add " + name.replace("_", " ").capitalize() + " (metarig)"
 147             op_type.bl_options = {'REGISTER', 'UNDO'}
 148             op_type.execute = make_metarig_add_execute(m)
 149
 150             metarig_ops[metarig_category].append((op_type, name))
 151
 152
 153 def create_menu_funcs():
 154     global menu_funcs
 155     for mop, name in metarig_ops[METARIG_DIR]:
 156         text = capwords(name.replace("_", " ")) + " (Meta-Rig)"
 157         menu_funcs += [make_metarig_menu_func(mop.bl_idname, text)]
 158
 159
 160 def create_armature_submenus(dic: dict | None = None):
 161     if dic is None:
 162         dic = metarigs
 163     global menu_funcs
 164     metarig_categories = list(dic.keys())
 165     metarig_categories.sort()
 166     for metarig_category in metarig_categories:
 167         # Create menu functions
 168         if metarig_category == "external":
 169             create_armature_submenus(dic=metarigs["external"])
 170             continue
 171         if metarig_category == METARIG_DIR:
 172             continue
 173
 174         armature_submenus.append(type('Class_' + metarig_category + '_submenu',
 175                                       (ArmatureSubMenu,), {}))
 176         armature_submenus[-1].bl_label = metarig_category + ' (submenu)'
 177         armature_submenus[-1].bl_idname = 'ARMATURE_MT_%s_class' % metarig_category
 178         armature_submenus[-1].operators = []
 179         menu_funcs += [make_submenu_func(armature_submenus[-1].bl_idname, metarig_category)]
 180
 181         for mop, name in metarig_ops[metarig_category]:
 182             arm_sub = next((e for e in armature_submenus
 183                             if e.bl_label == metarig_category + ' (submenu)'),
 184                            '')
 185             arm_sub.operators.append((mop.bl_idname, name,))
 186
 187
 188 def init_metarig_menu():
 189     get_internal_metarigs()
 190     create_metarig_ops()
 191     create_menu_funcs()
 192     create_armature_submenus()
 193
 194
 195 #################
 196 # Registering
 197
 198 def register():
 199     from bpy.utils import register_class
 200
 201     for cl in metarig_ops:
 202         for mop, name in metarig_ops[cl]:
 203             register_class(mop)
 204
 205     for arm_sub in armature_submenus:
 206         register_class(arm_sub)
 207
 208     for mf in menu_funcs:
 209         bpy.types.VIEW3D_MT_armature_add.append(mf)
 210
 211
 212 def unregister():
 213     from bpy.utils import unregister_class
 214
 215     for cl in metarig_ops:
 216         for mop, name in metarig_ops[cl]:
 217             unregister_class(mop)
 218
 219     for arm_sub in armature_submenus:
 220         unregister_class(arm_sub)
 221
 222     for mf in menu_funcs:
 223         bpy.types.VIEW3D_MT_armature_add.remove(mf)
 224
 225
 226 def get_external_metarigs(feature_module_names: list[str]):
 227     unregister()
 228
 229     # Clear and fill metarigs public variables
 230     metarigs.clear()
 231     get_internal_metarigs()
 232
 233     for module_name in feature_module_names:
 234         # noinspection PyBroadException
 235         try:
 236             base_dir, base_path = feature_set_list.get_dir_path(module_name, METARIG_DIR)
 237
 238             get_metarigs(metarigs['external'], base_dir, base_path)
 239         except Exception:
 240             print(f"Rigify Error: Could not load feature set '{module_name}' metarigs: "
 241                   f"exception occurred.\n")
 242             traceback.print_exc()
 243             print("")
 244             feature_set_list.mark_feature_set_exception(module_name)
 245             continue
 246
 247     metarig_ops.clear()
 248     armature_submenus.clear()
 249     menu_funcs.clear()
 250
 251     create_metarig_ops()
 252     create_menu_funcs()
 253     create_armature_submenus()
 254     register()