rigify/feature_set_list.py

   1 # SPDX-License-Identifier: GPL-2.0-or-later
   2
   3 from typing import List
   4
   5 import bpy
   6 from bpy.props import StringProperty
   7 import os
   8 import re
   9 import importlib
  10 import traceback
  11 from zipfile import ZipFile
  12 from shutil import rmtree
  13
  14 from . import feature_sets
  15
  16
  17 DEFAULT_NAME = 'rigify'
  18
  19 INSTALL_PATH = feature_sets._install_path()
  20 NAME_PREFIX = feature_sets.__name__.split('.')
  21
  22
  23 def get_install_path(*, create=False):
  24     if not os.path.exists(INSTALL_PATH):
  25         if create:
  26             os.makedirs(INSTALL_PATH, exist_ok=True)
  27         else:
  28             return None
  29
  30     return INSTALL_PATH
  31
  32
  33 def get_installed_modules_names() -> List[str]:
  34     """Return a list of module names of all feature sets in the file system."""
  35     features_path = get_install_path()
  36     if not features_path:
  37         return []
  38
  39     sets = []
  40
  41     for fs in os.listdir(features_path):
  42         if fs and fs[0] != '.' and fs != DEFAULT_NAME:
  43             fs_path = os.path.join(features_path, fs)
  44             if os.path.isdir(fs_path):
  45                 sets.append(fs)
  46
  47     return sets
  48
  49
  50 def get_enabled_modules_names() -> List[str]:
  51     """Return a list of module names of all enabled feature sets."""
  52     rigify_prefs = bpy.context.preferences.addons[__package__].preferences
  53     installed_module_names = get_installed_modules_names()
  54     rigify_feature_sets = rigify_prefs.rigify_feature_sets
  55
  56     enabled_module_names = { fs.module_name for fs in rigify_feature_sets if fs.enabled }
  57
  58     return [name for name in installed_module_names if name in enabled_module_names]
  59
  60
  61 def get_module(feature_set):
  62     return importlib.import_module('.'.join([*NAME_PREFIX, feature_set]))
  63
  64
  65 def get_module_safe(feature_set):
  66     try:
  67         return get_module(feature_set)
  68     except:
  69         return None
  70
  71
  72 def get_dir_path(feature_set, *extra_items):
  73     base_dir = os.path.join(INSTALL_PATH, feature_set, *extra_items)
  74     base_path = [*NAME_PREFIX, feature_set, *extra_items]
  75     return base_dir, base_path
  76
  77
  78 def get_info_dict(feature_set):
  79     module = get_module_safe(feature_set)
  80
  81     if module and hasattr(module, 'rigify_info'):
  82         data = module.rigify_info
  83         if isinstance(data, dict):
  84             return data
  85
  86     return {}
  87
  88
  89 def call_function_safe(module_name, func_name, args=[], kwargs={}):
  90     module = get_module_safe(module_name)
  91
  92     if module:
  93         func = getattr(module, func_name, None)
  94
  95         if callable(func):
  96             try:
  97                 return func(*args, **kwargs)
  98             except Exception:
  99                 print(f"Rigify Error: Could not call function '{func_name}' of feature set '{module_name}': exception occurred.\n")
 100                 traceback.print_exc()
 101                 print("")
 102
 103     return None
 104
 105
 106 def call_register_function(feature_set, register):
 107     call_function_safe(feature_set, 'register' if register else 'unregister')
 108
 109
 110 def get_ui_name(feature_set):
 111     # Try to get user-defined name
 112     info = get_info_dict(feature_set)
 113     if 'name' in info:
 114         return info['name']
 115
 116     # Default name based on directory
 117     name = re.sub(r'[_.-]', ' ', feature_set)
 118     name = re.sub(r'(?<=\d) (?=\d)', '.', name)
 119     return name.title()
 120
 121
 122 def feature_set_items(scene, context):
 123     """Get items for the Feature Set EnumProperty"""
 124     items = [
 125         ('all', 'All', 'All installed feature sets and rigs bundled with Rigify'),
 126         ('rigify', 'Rigify Built-in', 'Rigs bundled with Rigify'),
 127     ]
 128
 129     for fs in get_enabled_modules_names():
 130         ui_name = get_ui_name(fs)
 131         items.append((fs, ui_name, ui_name))
 132
 133     return items
 134
 135
 136 def verify_feature_set_archive(zipfile):
 137     """Verify that the zip file contains one root directory, and some required files."""
 138     dirname = None
 139     init_found = False
 140     data_found = False
 141
 142     for name in zipfile.namelist():
 143         parts = re.split(r'[/\\]', name)
 144
 145         if dirname is None:
 146             dirname = parts[0]
 147         elif dirname != parts[0]:
 148             dirname = None
 149             break
 150
 151         if len(parts) == 2 and parts[1] == '__init__.py':
 152             init_found = True
 153
 154         if len(parts) > 2 and parts[1] in {'rigs', 'metarigs'} and parts[-1] == '__init__.py':
 155             data_found = True
 156
 157     return dirname, init_found, data_found
 158
 159
 160 class DATA_OT_rigify_add_feature_set(bpy.types.Operator):
 161     bl_idname = "wm.rigify_add_feature_set"
 162     bl_label = "Add External Feature Set"
 163     bl_description = "Add external feature set (rigs, metarigs, ui templates)"
 164     bl_options = {"REGISTER", "UNDO", "INTERNAL"}
 165
 166     filter_glob: StringProperty(default="*.zip", options={'HIDDEN'})
 167     filepath: StringProperty(maxlen=1024, subtype='FILE_PATH', options={'HIDDEN', 'SKIP_SAVE'})
 168
 169     @classmethod
 170     def poll(cls, context):
 171         return True
 172
 173     def invoke(self, context, event):
 174         context.window_manager.fileselect_add(self)
 175         return {'RUNNING_MODAL'}
 176
 177     def execute(self, context):
 178         addon_prefs = context.preferences.addons[__package__].preferences
 179
 180         rigify_config_path = get_install_path(create=True)
 181
 182         with ZipFile(bpy.path.abspath(self.filepath), 'r') as zip_archive:
 183             base_dirname, init_found, data_found = verify_feature_set_archive(zip_archive)
 184
 185             if not base_dirname:
 186                 self.report({'ERROR'}, "The feature set archive must contain one base directory.")
 187                 return {'CANCELLED'}
 188
 189             # Patch up some invalid characters to allow using 'Download ZIP' on GitHub.
 190             fixed_dirname = re.sub(r'[.-]', '_', base_dirname)
 191
 192             if not re.fullmatch(r'[a-zA-Z][a-zA-Z_0-9]*', fixed_dirname):
 193                 self.report({'ERROR'}, "The feature set archive base directory name is not a valid identifier: '%s'." % (base_dirname))
 194                 return {'CANCELLED'}
 195
 196             if fixed_dirname == DEFAULT_NAME:
 197                 self.report({'ERROR'}, "The '%s' name is not allowed for feature sets." % (DEFAULT_NAME))
 198                 return {'CANCELLED'}
 199
 200             if not init_found or not data_found:
 201                 self.report({'ERROR'}, "The feature set archive has no rigs or metarigs, or is missing __init__.py.")
 202                 return {'CANCELLED'}
 203
 204             base_dir = os.path.join(rigify_config_path, base_dirname)
 205             fixed_dir = os.path.join(rigify_config_path, fixed_dirname)
 206
 207             for path, name in [(base_dir, base_dirname), (fixed_dir, fixed_dirname)]:
 208                 if os.path.exists(path):
 209                     self.report({'ERROR'}, "Feature set directory already exists: '%s'." % (name))
 210                     return {'CANCELLED'}
 211
 212             # Unpack the validated archive and fix the directory name if necessary
 213             zip_archive.extractall(rigify_config_path)
 214
 215             if base_dir != fixed_dir:
 216                 os.rename(base_dir, fixed_dir)
 217
 218             # Call the register callback of the new set
 219             call_register_function(fixed_dirname, True)
 220
 221             addon_prefs.update_external_rigs()
 222
 223             addon_prefs.active_feature_set_index = len(addon_prefs.rigify_feature_sets)-1
 224
 225         return {'FINISHED'}
 226
 227
 228 class DATA_OT_rigify_remove_feature_set(bpy.types.Operator):
 229     bl_idname = "wm.rigify_remove_feature_set"
 230     bl_label = "Remove External Feature Set"
 231     bl_description = "Remove external feature set (rigs, metarigs, ui templates)"
 232     bl_options = {"REGISTER", "UNDO", "INTERNAL"}
 233
 234     @classmethod
 235     def poll(cls, context):
 236         return True
 237
 238     def invoke(self, context, event):
 239         return context.window_manager.invoke_confirm(self, event)
 240
 241     def execute(self, context):
 242         addon_prefs = context.preferences.addons[__package__].preferences
 243         feature_sets = addon_prefs.rigify_feature_sets
 244         active_idx = addon_prefs.active_feature_set_index
 245         active_fs = feature_sets[active_idx]
 246
 247         # Call the unregister callback of the set being removed.
 248         if active_fs.enabled:
 249             call_register_function(active_fs.module_name, register=False)
 250
 251         # Remove the feature set's folder from the file system.
 252         rigify_config_path = get_install_path()
 253         if rigify_config_path:
 254             set_path = os.path.join(rigify_config_path, active_fs.module_name)
 255             if os.path.exists(set_path):
 256                 rmtree(set_path)
 257
 258         # Remove the feature set's entry from the addon preferences.
 259         feature_sets.remove(active_idx)
 260
 261         # Remove the feature set's entries from the metarigs and rig types.
 262         addon_prefs.update_external_rigs()
 263
 264         # Update active index.
 265         addon_prefs.active_feature_set_index -= 1
 266
 267         return {'FINISHED'}
 268
 269
 270 def register():
 271     bpy.utils.register_class(DATA_OT_rigify_add_feature_set)
 272     bpy.utils.register_class(DATA_OT_rigify_remove_feature_set)
 273
 274 def unregister():
 275     bpy.utils.unregister_class(DATA_OT_rigify_add_feature_set)
 276     bpy.utils.unregister_class(DATA_OT_rigify_remove_feature_set)