rigify/feature_set_list.py

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