Rigify: support register and unregister functions in feature sets.

Allow feature sets to register their own operators, in case
e.g. they want to add a button to the metarig UI panel.

2 files changed, 38 insertions, 0 deletions

diff --git a/rigify/__init__.py b/rigify/__init__.py
index 00777a78..6ae597e9 100644
--- a/rigify/__init__.py
+++ b/rigify/__init__.py
@@ -195,6 +195,14 @@ class RigifyPreferences(AddonPreferences):
 
             register()
 
+    def register_feature_sets(self, register):
+        """Call register or unregister of external feature sets"""
+        if self.legacy_mode:
+            return
+
+        for set_name in feature_set_list.get_installed_list():
+            feature_set_list.call_register_function(set_name, register)
+
     def update_external_rigs(self, force=False):
         """Get external feature sets"""
         if self.legacy_mode:
@@ -559,6 +567,7 @@ def register():
     if legacy_loaded or bpy.context.preferences.addons['rigify'].preferences.legacy_mode:
         bpy.context.preferences.addons['rigify'].preferences.legacy_mode = True
 
+    bpy.context.preferences.addons['rigify'].preferences.register_feature_sets(True)
     bpy.context.preferences.addons['rigify'].preferences.update_external_rigs()
 
     # Add rig parameters
@@ -589,6 +598,8 @@ def register_rig_parameters():
 def unregister():
     from bpy.utils import unregister_class
 
+    bpy.context.preferences.addons['rigify'].preferences.register_feature_sets(False)
+
     # Properties on PoseBones and Armature.
     del bpy.types.PoseBone.rigify_type
     del bpy.types.PoseBone.rigify_parameters
diff --git a/rigify/feature_set_list.py b/rigify/feature_set_list.py
index 524363ad..5a2d928b 100644
--- a/rigify/feature_set_list.py
+++ b/rigify/feature_set_list.py
@@ -87,6 +87,27 @@ def get_info_dict(feature_set):
     return {}
 
 
+def call_function_safe(feature_set, name, args=[], kwargs={}):
+    module = get_module_safe(feature_set)
+
+    if module:
+        func = getattr(module, name, None)
+
+        if callable(func):
+            try:
+                return func(*args, **kwargs)
+            except Exception:
+                print("Rigify Error: Could not call function '%s' of feature set '%s': exception occurred.\n" % (name, feature_set))
+                traceback.print_exc()
+                print("")
+
+    return None
+
+
+def call_register_function(feature_set, register):
+    call_function_safe(feature_set, 'register' if register else 'unregister')
+
+
 def get_ui_name(feature_set):
     # Try to get user-defined name
     info = get_info_dict(feature_set)
@@ -191,6 +212,9 @@ class DATA_OT_rigify_add_feature_set(bpy.types.Operator):
             if base_dir != fixed_dir:
                 os.rename(base_dir, fixed_dir)
 
+            # Call the register callback of the new set
+            call_register_function(fixed_dirname, True)
+
         addon_prefs.machin = bpy.props.EnumProperty(items=(('a',)*3, ('b',)*3, ('c',)*3),)
 
         addon_prefs.update_external_rigs()
@@ -215,6 +239,9 @@ class DATA_OT_rigify_remove_feature_set(bpy.types.Operator):
     def execute(self, context):
         addon_prefs = context.preferences.addons[__package__].preferences
 
+        # Call the unregister callback of the set being removed
+        call_register_function(self.featureset, False)
+
         rigify_config_path = get_install_path()
         if rigify_config_path:
             set_path = os.path.join(rigify_config_path, self.featureset)