rigify/metarig_menu.py

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