rigify/feature_set_list.py

   1 # SPDX-FileCopyrightText: 2019-2022 Blender Foundation
   2 #
   3 # SPDX-License-Identifier: GPL-2.0-or-later
   4
   5 from typing import TYPE_CHECKING, List, Sequence, Optional
   6
   7 import bpy
   8 from bpy.props import StringProperty
   9 import os
  10 import re
  11 import importlib
  12 import traceback
  13 from zipfile import ZipFile
  14 from shutil import rmtree
  15
  16 from . import feature_sets
  17
  18 if TYPE_CHECKING:
  19     from . import RigifyFeatureSets
  20
  21
  22 DEFAULT_NAME = 'rigify'
  23
  24 # noinspection PyProtectedMember
  25 INSTALL_PATH = feature_sets._install_path()
  26 NAME_PREFIX = feature_sets.__name__.split('.')
  27
  28 # noinspection SpellCheckingInspection
  29 PROMOTED_FEATURE_SETS = [
  30     {
  31         "name": "CloudRig",
  32         "author": "Demeter Dzadik",
  33         "description": "Feature set developed by the Blender Animation Studio",
  34         "doc_url": "https://gitlab.com/blender/CloudRig/-/wikis/",
  35         "link": "https://gitlab.com/blender/CloudRig/",
  36     },
  37     {
  38         "name": "Experimental Rigs by Alexander Gavrilov",
  39         "author": "Alexander Gavrilov",
  40         "description":
  41             "Experimental and/or niche rigs made by a Rigify maintainer.\n"
  42             "Includes a BlenRig-like spine, Body IK (knee & elbow IK), jiggles, skin transforms, etc.",
  43         "link": "https://github.com/angavrilov/angavrilov-rigs",
  44     },
  45     {
  46         "name": "Cessen's Rigify Extensions",
  47         "author": "Nathan Vegdahl",
  48         "description": "Collection of original legacy Rigify rigs minimally ported to the modern Rigify",
  49         "warning": "This feature set is maintained at the bare minimal level",
  50         "link": "https://github.com/cessen/cessen_rigify_ext",
  51     }
  52 ]
  53
  54
  55 def get_install_path(*, create=False):
  56     if not os.path.exists(INSTALL_PATH):
  57         if create:
  58             os.makedirs(INSTALL_PATH, exist_ok=True)
  59         else:
  60             return None
  61
  62     return INSTALL_PATH
  63
  64
  65 def get_installed_modules_names() -> List[str]:
  66     """Return a list of module names of all feature sets in the file system."""
  67     features_path = get_install_path()
  68     if not features_path:
  69         return []
  70
  71     sets = []
  72
  73     for fs in os.listdir(features_path):
  74         if fs and fs[0] != '.' and fs != DEFAULT_NAME:
  75             fs_path = os.path.join(features_path, fs)
  76             if os.path.isdir(fs_path):
  77                 sets.append(fs)
  78
  79     return sets
  80
  81
  82 def get_prefs_feature_sets() -> Sequence['RigifyFeatureSets']:
  83     from . import RigifyPreferences
  84     return RigifyPreferences.get_instance().rigify_feature_sets
  85
  86
  87 def get_enabled_modules_names() -> List[str]:
  88     """Return a list of module names of all enabled feature sets."""
  89     installed_module_names = get_installed_modules_names()
  90     rigify_feature_sets = get_prefs_feature_sets()
  91
  92     enabled_module_names = {fs.module_name for fs in rigify_feature_sets if fs.enabled}
  93
  94     return [name for name in installed_module_names if name in enabled_module_names]
  95
  96
  97 def find_module_name_by_link(link: str) -> str | None:
  98     """Returns the name of the feature set module that is associated with the specified url."""
  99     if not link:
 100         return None
 101
 102     for fs in get_prefs_feature_sets():
 103         if fs.link == link:
 104             return fs.module_name or None
 105
 106     return None
 107
 108
 109 def mark_feature_set_exception(module_name: str):
 110     if module_name:
 111         for fs in get_prefs_feature_sets():
 112             if fs.module_name == module_name:
 113                 fs.has_exceptions = True
 114
 115
 116 def get_module(feature_set: str):
 117     return importlib.import_module('.'.join([*NAME_PREFIX, feature_set]))
 118
 119
 120 def get_module_safe(feature_set: str):
 121     # noinspection PyBroadException
 122     try:
 123         return get_module(feature_set)
 124     except:  # noqa: E722
 125         return None
 126
 127
 128 def get_module_by_link_safe(link: str):
 129     if module_name := find_module_name_by_link(link):
 130         return get_module_safe(module_name)
 131
 132
 133 def get_dir_path(feature_set: str, *extra_items: list[str]):
 134     base_dir = os.path.join(INSTALL_PATH, feature_set, *extra_items)
 135     base_path = [*NAME_PREFIX, feature_set, *extra_items]
 136     return base_dir, base_path
 137
 138
 139 def get_info_dict(feature_set: str):
 140     module = get_module_safe(feature_set)
 141
 142     if module and hasattr(module, 'rigify_info'):
 143         data = module.rigify_info
 144         if isinstance(data, dict):
 145             return data
 146
 147     return {}
 148
 149
 150 def call_function_safe(module_name: str, func_name: str,
 151                        args: Optional[list] = None, kwargs: Optional[dict] = None,
 152                        mark_error=False):
 153     module = get_module_safe(module_name)
 154
 155     if module:
 156         func = getattr(module, func_name, None)
 157
 158         if callable(func):
 159             # noinspection PyBroadException
 160             try:
 161                 return func(*(args or []), **(kwargs or {}))
 162             except Exception:
 163                 print(f"Rigify Error: Could not call function '{func_name}' of feature set "
 164                       f"'{module_name}': exception occurred.\n")
 165                 traceback.print_exc()
 166                 print("")
 167
 168                 if mark_error:
 169                     mark_feature_set_exception(module_name)
 170
 171     return None
 172
 173
 174 def call_register_function(feature_set: str, do_register: bool):
 175     call_function_safe(feature_set, 'register' if do_register else 'unregister', mark_error=do_register)
 176
 177
 178 def get_ui_name(feature_set: str):
 179     # Try to get user-defined name
 180     info = get_info_dict(feature_set)
 181     if 'name' in info:
 182         return info['name']
 183
 184     # Default name based on directory
 185     name = re.sub(r'[_.-]', ' ', feature_set)
 186     name = re.sub(r'(?<=\d) (?=\d)', '.', name)
 187     return name.title()
 188
 189
 190 def feature_set_items(_scene, _context):
 191     """Get items for the Feature Set EnumProperty"""
 192     items = [
 193         ('all', 'All', 'All installed feature sets and rigs bundled with Rigify'),
 194         ('rigify', 'Rigify Built-in', 'Rigs bundled with Rigify'),
 195     ]
 196
 197     for fs in get_enabled_modules_names():
 198         ui_name = get_ui_name(fs)
 199         items.append((fs, ui_name, ui_name))
 200
 201     return items
 202
 203
 204 def verify_feature_set_archive(zipfile):
 205     """Verify that the zip file contains one root directory, and some required files."""
 206     dirname = None
 207     init_found = False
 208     data_found = False
 209
 210     for name in zipfile.namelist():
 211         parts = re.split(r'[/\\]', name)
 212
 213         if dirname is None:
 214             dirname = parts[0]
 215         elif dirname != parts[0]:
 216             dirname = None
 217             break
 218
 219         if len(parts) == 2 and parts[1] == '__init__.py':
 220             init_found = True
 221
 222         if len(parts) > 2 and parts[1] in {'rigs', 'metarigs'} and parts[-1] == '__init__.py':
 223             data_found = True
 224
 225     return dirname, init_found, data_found
 226
 227
 228 # noinspection PyPep8Naming
 229 class DATA_OT_rigify_add_feature_set(bpy.types.Operator):
 230     bl_idname = "wm.rigify_add_feature_set"
 231     bl_label = "Add External Feature Set"
 232     bl_description = "Add external feature set (rigs, metarigs, ui templates)"
 233     bl_options = {"REGISTER", "UNDO", "INTERNAL"}
 234
 235     filter_glob: StringProperty(default="*.zip", options={'HIDDEN'})
 236     filepath: StringProperty(maxlen=1024, subtype='FILE_PATH', options={'HIDDEN', 'SKIP_SAVE'})
 237
 238     @classmethod
 239     def poll(cls, context):
 240         return True
 241
 242     def invoke(self, context, event):
 243         context.window_manager.fileselect_add(self)
 244         return {'RUNNING_MODAL'}
 245
 246     def execute(self, context):
 247         from . import RigifyPreferences
 248         addon_prefs = RigifyPreferences.get_instance()
 249
 250         rigify_config_path = get_install_path(create=True)
 251
 252         with ZipFile(bpy.path.abspath(self.filepath), 'r') as zip_archive:
 253             base_dirname, init_found, data_found = verify_feature_set_archive(zip_archive)
 254
 255             if not base_dirname:
 256                 self.report({'ERROR'}, "The feature set archive must contain one base directory.")
 257                 return {'CANCELLED'}
 258
 259             # Patch up some invalid characters to allow using 'Download ZIP' on GitHub.
 260             fixed_dirname = re.sub(r'[.-]', '_', base_dirname)
 261
 262             if not re.fullmatch(r'[a-zA-Z][a-zA-Z_0-9]*', fixed_dirname):
 263                 self.report({'ERROR'},
 264                             f"The feature set archive base directory name is not a valid "
 265                             f"identifier: '{base_dirname}'.")
 266                 return {'CANCELLED'}
 267
 268             if fixed_dirname == DEFAULT_NAME:
 269                 self.report(
 270                     {'ERROR'}, f"The '{DEFAULT_NAME}' name is not allowed for feature sets.")
 271                 return {'CANCELLED'}
 272
 273             if not init_found or not data_found:
 274                 self.report(
 275                     {'ERROR'},
 276                     "The feature set archive has no rigs or metarigs, or is missing __init__.py.")
 277                 return {'CANCELLED'}
 278
 279             base_dir = os.path.join(rigify_config_path, base_dirname)
 280             fixed_dir = os.path.join(rigify_config_path, fixed_dirname)
 281
 282             for path, name in [(base_dir, base_dirname), (fixed_dir, fixed_dirname)]:
 283                 if os.path.exists(path):
 284                     self.report({'ERROR'}, f"Feature set directory already exists: '{name}'.")
 285                     return {'CANCELLED'}
 286
 287             # Unpack the validated archive and fix the directory name if necessary
 288             zip_archive.extractall(rigify_config_path)
 289
 290             if base_dir != fixed_dir:
 291                 os.rename(base_dir, fixed_dir)
 292
 293             # Call the register callback of the new set
 294             addon_prefs.refresh_installed_feature_sets()
 295
 296             call_register_function(fixed_dirname, True)
 297
 298             addon_prefs.update_external_rigs()
 299
 300             # Select the new entry
 301             for i, fs in enumerate(addon_prefs.rigify_feature_sets):
 302                 if fs.module_name == fixed_dirname:
 303                     addon_prefs.active_feature_set_index = i
 304                     break
 305
 306         return {'FINISHED'}
 307
 308
 309 # noinspection PyPep8Naming
 310 class DATA_OT_rigify_remove_feature_set(bpy.types.Operator):
 311     bl_idname = "wm.rigify_remove_feature_set"
 312     bl_label = "Remove External Feature Set"
 313     bl_description = "Remove external feature set (rigs, metarigs, ui templates)"
 314     bl_options = {"REGISTER", "UNDO", "INTERNAL"}
 315
 316     @classmethod
 317     def poll(cls, context):
 318         return True
 319
 320     def invoke(self, context, event):
 321         return context.window_manager.invoke_confirm(self, event)
 322
 323     def execute(self, context):
 324         from . import RigifyPreferences
 325         addon_prefs = RigifyPreferences.get_instance()
 326         feature_set_list = addon_prefs.rigify_feature_sets
 327         active_idx = addon_prefs.active_feature_set_index
 328         active_fs: 'RigifyFeatureSets' = feature_set_list[active_idx]
 329
 330         # Call the 'unregister' callback of the set being removed.
 331         if active_fs.enabled:
 332             call_register_function(active_fs.module_name, do_register=False)
 333
 334         # Remove the feature set's folder from the file system.
 335         rigify_config_path = get_install_path()
 336         if rigify_config_path:
 337             set_path = os.path.join(rigify_config_path, active_fs.module_name)
 338             if os.path.exists(set_path):
 339                 rmtree(set_path)
 340
 341         # Remove the feature set's entry from the addon preferences.
 342         feature_set_list.remove(active_idx)
 343
 344         # Remove the feature set's entries from the metarigs and rig types.
 345         addon_prefs.refresh_installed_feature_sets()
 346         addon_prefs.update_external_rigs()
 347
 348         # Update active index.
 349         addon_prefs.active_feature_set_index -= 1
 350
 351         return {'FINISHED'}
 352
 353
 354 def register():
 355     bpy.utils.register_class(DATA_OT_rigify_add_feature_set)
 356     bpy.utils.register_class(DATA_OT_rigify_remove_feature_set)
 357
 358
 359 def unregister():
 360     bpy.utils.unregister_class(DATA_OT_rigify_add_feature_set)
 361     bpy.utils.unregister_class(DATA_OT_rigify_remove_feature_set)