Support declaring new callback methods via stage decorators.

Making generic rig 'template' classes that are intended to be
subclassed in specific rigs involves splitting operations done
in each stage into multiple methods that can be overridden
separately. The main callback thus ends up simply calling a
sequence of other methods.

To make such code cleaner it's better to allow registering
those methods as new callbacks that would be automatically
called by the system. This can be done via decorators.

A new metaclass used for all rig and generate plugin classes
builds and validates a table of all decorated methods, and
allows calling them all together with the main callback.

Author: angavrilov

base_generate.py

@@ -171,7 +171,7 @@ def __run_object_stage(self, method_name):
         self.stage = method_name
 
         for rig in [*self.rig_list, *self.plugin_list]:
-            getattr(rig, method_name)()
+            rig.rigify_invoke_stage(method_name)
 
             assert(self.context.active_object == self.obj)
             assert(self.obj.mode == 'OBJECT')
@@ -186,7 +186,7 @@ def __run_edit_stage(self, method_name):
         self.stage = method_name
 
         for rig in [*self.rig_list, *self.plugin_list]:
-            getattr(rig, method_name)()
+            rig.rigify_invoke_stage(method_name)
 
             assert(self.context.active_object == self.obj)
             assert(self.obj.mode == 'EDIT')
@@ -221,15 +221,15 @@ def invoke_generate_bones(self):
         self.stage = 'generate_bones'
 
         for rig in self.rig_list:
-            rig.generate_bones()
+            rig.rigify_invoke_stage('generate_bones')
 
             assert(self.context.active_object == self.obj)
             assert(self.obj.mode == 'EDIT')
 
             self.__auto_register_bones(self.obj.data.edit_bones, rig)
 
         for plugin in self.plugin_list:
-            plugin.generate_bones()
+            plugin.rigify_invoke_stage('generate_bones')
 
             assert(self.context.active_object == self.obj)
             assert(self.obj.mode == 'EDIT')

base_rig.py

@@ -24,6 +24,7 @@
 
 from .utils.bones import BoneDict, BoneUtilityMixin
 from .utils.mechanism import MechanismUtilityMixin
+from .utils.metaclass import BaseStagedClass
 
 # Only export certain symbols via 'from base_rig import *'
 __all__ = ['BaseRig', 'RigUtility']
@@ -32,7 +33,7 @@
 # Base Rig
 #=============================================
 
-class GenerateCallbackMixin(object):
+class GenerateCallbackMixin(BaseStagedClass):
     """
     Standard set of callback methods to redefine.
     Shared between BaseRig and GeneratorPlugin.
@@ -43,7 +44,27 @@ class GenerateCallbackMixin(object):
     Switching modes is not allowed in rigs for performance
     reasons. Place code in the appropriate callbacks to use
     the mode set by the main engine.
+
+    Before each callback, all other methods decorated with
+    @stage_<method_name> are called, for instance:
+
+    @stage_generate_bones
+    def foo(self):
+        print('first')
+
+    def generate_bones(self):
+        print('second')
+
+    Will print 'first', then 'second'. However, the order
+    in which different @stage_generate_bones methods in the
+    same rig will be called is not specified.
+
+    When overriding such methods in a subclass the appropriate
+    decorator should be repeated for code clarity reasons;
+    a warning is printed if this is not done.
     """
+    DEFINE_STAGES = True
+
     def initialize(self):
         """
         Initialize processing after all rig classes are constructed.
@@ -203,3 +224,11 @@ def register_new_bone(self, new_name, old_name=None):
         self.owner.register_new_bone(new_name, old_name)
 
 
+#=============================================
+# Rig Stage Decorators
+#=============================================
+
+# Generate and export @stage_... decorators for all valid stages
+for name, decorator in GenerateCallbackMixin.make_stage_decorators():
+    globals()[name] = decorator
+    __all__.append(name)

utils/metaclass.py

@@ -20,13 +20,118 @@
 
 import collections
 
+from types import FunctionType
+
+
+#=============================================
+# Class With Stages
+#=============================================
+
+
+def rigify_stage(stage):
+    """Decorates the method with the specified stage."""
+    def process(method):
+        if not isinstance(method, FunctionType):
+            raise ValueError("Stage decorator must be applied to a method definition")
+        method._rigify_stage = stage
+        return method
+    return process
+
+
+class StagedMetaclass(type):
+    """
+    Metaclass for rigs that manages assignment of methods to stages via @stage_* decorators.
+
+    Using 'DEFINE_STAGES = True' inside the class definition will register all non-system
+    method names from that definition as valid stages. After that, subclasses can
+    register methods to those stages, to be called via rigify_invoke_stage.
+    """
+    def __new__(metacls, class_name, bases, namespace, **kwds):
+        staged_bases = [base for base in bases if isinstance(base, StagedMetaclass)]
+
+        # Compute the set of inherited stages
+        stages = set().union(*[base._rigify_stages for base in staged_bases])
+
+        # Add methods from current class if requested
+        if 'DEFINE_STAGES' in namespace:
+            del namespace['DEFINE_STAGES']
+
+            for name, item in namespace.items():
+                if name[0] != '_' and isinstance(item, FunctionType):
+                    stages.add(name)
+
+        # Create the class
+        result = type.__new__(metacls, class_name, bases, dict(namespace))
+
+        # Compute the inherited stage to method mapping
+        stage_map = collections.defaultdict(collections.OrderedDict)
+        method_map = {}
+
+        for base in staged_bases:
+            for stage_name, methods in base._rigify_stage_map.items():
+                for method_name, method_class in methods.items():
+                    # Check consistency of inherited stage assignment to methods
+                    if method_name in method_map:
+                        if method_map[method_name] != stage_name:
+                            print("RIGIFY CLASS %s (%s): method '%s' has inherited both @stage_%s and @stage_%s\n" %
+                                  (class_name, result.__module__, method_name, method_map[method_name], stage_name))
+                    else:
+                        method_map[method_name] = stage_name
+
+                    stage_map[stage_name][method_name] = method_class
+
+        # Scan newly defined methods for stage decorations
+        for method_name, item in namespace.items():
+            if isinstance(item, FunctionType):
+                stage = getattr(item, '_rigify_stage', None)
+
+                # Ensure that decorators aren't lost when redefining methods
+                if method_name in method_map:
+                    if not stage:
+                        stage = method_map[method_name]
+                        print("RIGIFY CLASS %s (%s): missing stage decorator on method '%s' (should be @stage_%s)" %
+                              (class_name, result.__module__, method_name, stage))
+                    # Check that the method is assigned to only one stage
+                    elif stage != method_map[method_name]:
+                        print("RIGIFY CLASS %s (%s): method '%s' has decorator @stage_%s, but inherited base has @stage_%s" %
+                              (class_name, result.__module__, method_name, stage, method_map[method_name]))
+
+                # Assign the method to the stage, verifying that it's valid
+                if stage:
+                    if stage not in stages:
+                        raise ValueError("Invalid stage name '%s' for method '%s' in class %s (%s)" %
+                                         (stage, method_name, class_name, result.__module__))
+                    else:
+                        stage_map[stage][method_name] = result
+
+        result._rigify_stages = frozenset(stages)
+        result._rigify_stage_map = stage_map
+
+        return result
+
+    def make_stage_decorators(self):
+        return [('stage_'+name, rigify_stage(name)) for name in self._rigify_stages]
+
+
+class BaseStagedClass(object, metaclass=StagedMetaclass):
+    def rigify_invoke_stage(self, stage):
+        """Call all methods decorated with the given stage, followed by the callback."""
+        cls = self.__class__
+        assert(isinstance(cls, StagedMetaclass))
+        assert(stage in cls._rigify_stages)
+
+        for method_name in cls._rigify_stage_map[stage]:
+            getattr(self, method_name)()
+
+        getattr(self, stage)()
+
 
 #=============================================
 # Per-owner singleton class
 #=============================================
 
 
-class SingletonPluginMetaclass(type):
+class SingletonPluginMetaclass(StagedMetaclass):
     """Metaclass for maintaining one instance per owner object per constructor arg set."""
     def __call__(cls, owner, *constructor_args):
         key = (cls, *constructor_args)