Make rig_ui_template modular: allow individual rigs to add code.

Allow rig classes to return new module imports, utility functions
and operators, and classes to register. This is necessary to make
script generation truly modular.

In order to avoid code duplication, rigs return lists of strings,
which are merged and filtered for duplicates before insertion into
the script. In this way, multiple rigs can easily share utility
code by simply ensuring that they include exactly the same string.

As an example, existing super_limb rigs are changed to use this
feature. For now arm and leg utilities are also included by default
to provide backward compatibility to legacy rigs.

Author: angavrilov

README

@@ -240,15 +240,39 @@ Implementation classes are shown in the metarig samples list and generate a samp
 
 GENERATING A PYTHON UI
 ----------------------
-The generate() method can also, optionally, return python code as a single
-string.  This python code is added to the "rig properties" panel that gets
-auto-generated along with the rig.  This is useful for exposing things like
-IK/FK switches in a nice way to the animator.
+The generate() method can also, optionally, return a dictionary with data
+used for generating the python UI. The dictionary may have the following
+keys (all optional):
 
-The string must be returned in a list, e.g.:
+  'script': python code as list of strings.
 
-return ["my python code"]
+    This python code is added to the "rig properties" panel that gets
+    auto-generated along with the rig.  This is useful for exposing things like
+    IK/FK switches in a nice way to the animator.
 
-The reason it needs to be put in a list is to leave room for expanding the API
-in the future, for returning additional information.
+  'utilities': python code as list of strings.
 
+    This code is added to the utility function/class area of the auto-generated
+    python script, and should contain functions, classes and operators for use
+    in the 'script' code and other utilities.
+
+  'imports': python code as list of strings.
+
+    These lines are added to the import section of the python script.
+
+  'register': list of class names.
+
+    These are used in generating the register() and unregister() functions.
+
+The 'utilities', 'imports' and 'register' lists returned by all rigs are concatenated
+and filtered for duplicates before use, in order to allow easy sharing of utility
+functions and classes: all that is needed is to ensure that shared code is included
+in the lists as exactly the same string by all users.
+
+The python code strings from the lists are implicitly terminated with a newline.
+
+Some rigs return a list with a single string, e.g.:
+
+    return ["my python code"]
+
+This is a deprecated format equivalent to {'script': ["my python code"]}.

generate.py

@@ -24,6 +24,7 @@
 import traceback
 import sys
 from rna_prop_ui import rna_idprop_ui_prop_get
+from collections import OrderedDict
 
 from .utils import MetarigError, new_bone, get_rig_type
 from .utils import ORG_PREFIX, MCH_PREFIX, DEF_PREFIX, WGT_PREFIX, ROOT_NAME, make_original_name
@@ -32,7 +33,7 @@
 from .utils import random_id
 from .utils import copy_attributes
 from .utils import gamma_correct
-from .rig_ui_template import UI_SLIDERS, layers_ui, UI_REGISTER
+from .rig_ui_template import UI_IMPORTS, UI_BASE_UTILITIES, UI_UTILITIES, UI_SLIDERS, layers_ui, UI_REGISTER
 
 
 RIG_MODULE = "rigs"
@@ -341,14 +342,26 @@ def generate_rig(context, metarig):
 
         # Generate all the rigs.
         ui_scripts = []
+        ui_imports = UI_IMPORTS.copy()
+        ui_utilities = UI_UTILITIES.copy()
+        ui_register = UI_REGISTER.copy()
         for rig in rigs:
             # Go into editmode in the rig armature
             bpy.ops.object.mode_set(mode='OBJECT')
             context.scene.objects.active = obj
             obj.select = True
             bpy.ops.object.mode_set(mode='EDIT')
             scripts = rig.generate()
-            if scripts is not None:
+            if isinstance(scripts, dict):
+                if 'script' in scripts:
+                    ui_scripts += scripts['script']
+                if 'imports' in scripts:
+                    ui_imports += scripts['imports']
+                if 'utilities' in scripts:
+                    ui_utilities += scripts['utilities']
+                if 'register' in scripts:
+                    ui_register += scripts['register']
+            elif scripts is not None:
                 ui_scripts += [scripts[0]]
         t.tick("Generate rigs: ")
 
@@ -496,11 +509,23 @@ def generate_rig(context, metarig):
 
     id_store.rigify_rig_ui = script.name
 
-    script.write(UI_SLIDERS % rig_id)
+    for s in OrderedDict.fromkeys(ui_imports):
+        script.write(s + "\n")
+    script.write(UI_BASE_UTILITIES % rig_id)
+    for s in OrderedDict.fromkeys(ui_utilities):
+        script.write(s + "\n")
+    script.write(UI_SLIDERS)
     for s in ui_scripts:
         script.write("\n        " + s.replace("\n", "\n        ") + "\n")
     script.write(layers_ui(vis_layers, layer_layout))
-    script.write(UI_REGISTER)
+    script.write("\ndef register():\n")
+    ui_register = OrderedDict.fromkeys(ui_register)
+    for s in ui_register:
+        script.write("    bpy.utils.register_class("+s+");\n")
+    script.write("\ndef unregister():\n")
+    for s in ui_register:
+        script.write("    bpy.utils.unregister_class("+s+");\n")
+    script.write("\nregister()\n")
     script.use_module = True
 
     # Run UI script

rig_ui_template.py

@@ -18,11 +18,14 @@
 
 # <pep8 compliant>
 
-UI_SLIDERS = '''
-import bpy
-from mathutils import Matrix, Vector
-from math import acos, pi, radians
-
+UI_IMPORTS = [
+    'import bpy',
+    'import math',
+    'from math import pi',
+    'from mathutils import Euler, Matrix, Quaternion, Vector',
+]
+
+UI_BASE_UTILITIES = '''
 rig_id = "%s"
 
 
@@ -54,7 +57,7 @@ def rotation_difference(mat1, mat2):
     """
     q1 = mat1.to_quaternion()
     q2 = mat2.to_quaternion()
-    angle = acos(min(1,max(-1,q1.dot(q2)))) * 2
+    angle = math.acos(min(1,max(-1,q1.dot(q2)))) * 2
     if angle > pi:
         angle = -angle + (2*pi)
     return angle
@@ -295,6 +298,22 @@ def set_pole(pvi):
     if ang1 < ang2:
         set_pole(pv1)
 
+##########
+## Misc ##
+##########
+
+def parse_bone_names(names_string):
+    if names_string[0] == '[' and names_string[-1] == ']':
+        return eval(names_string)
+    else:
+        return names_string
+
+'''
+
+UTILITIES_FUNC_ARM_FKIK = ['''
+######################
+## IK Arm functions ##
+######################
 
 def fk2ik_arm(obj, fk, ik):
     """ Matches the fk bones in an arm rig to the ik bones.
@@ -390,6 +409,12 @@ def ik2fk_arm(obj, fk, ik):
         match_pose_scale(uarmi, uarm)
         # Rotation Correction
         correct_rotation(uarmi, uarm)
+''']
+
+UTILITIES_FUNC_LEG_FKIK = ['''
+######################
+## IK Leg functions ##
+######################
 
 def fk2ik_leg(obj, fk, ik):
     """ Matches the fk bones in a leg rig to the ik bones.
@@ -518,18 +543,13 @@ def ik2fk_leg(obj, fk, ik):
 
         # Pole target position
         match_pole_target(thighi, shini, pole, thigh, (thighi.length + shini.length))
+''']
 
-
+UTILITIES_FUNC_POLE = ['''
 ################################
 ## IK Rotation-Pole functions ##
 ################################
 
-def parse_bone_names(names_string):
-    if names_string[0] == '[' and names_string[-1] == ']':
-        return eval(names_string)
-    else:
-        return names_string
-
 def rotPoleToggle(rig, limb_type, controls, ik_ctrl, fk_ctrl, parent, pole):
 
     rig_id = rig.data['rig_id']
@@ -588,10 +608,14 @@ def rotPoleToggle(rig, limb_type, controls, ik_ctrl, fk_ctrl, parent, pole):
             func2(**kwargs2)
 
             bpy.ops.pose.select_all(action='DESELECT')
+''']
 
-##############################
-## IK/FK snapping operators ##
-##############################
+REGISTER_OP_ARM_FKIK = ['Rigify_Arm_FK2IK', 'Rigify_Arm_IK2FK']
+
+UTILITIES_OP_ARM_FKIK = ['''
+##################################
+## IK/FK Arm snapping operators ##
+##################################
 
 class Rigify_Arm_FK2IK(bpy.types.Operator):
     """ Snaps an FK arm to an IK arm.
@@ -652,7 +676,14 @@ def execute(self, context):
         finally:
             context.user_preferences.edit.use_global_undo = use_global_undo
         return {'FINISHED'}
+''']
+
+REGISTER_OP_LEG_FKIK = ['Rigify_Leg_FK2IK', 'Rigify_Leg_IK2FK']
 
+UTILITIES_OP_LEG_FKIK = ['''
+##################################
+## IK/FK Leg snapping operators ##
+##################################
 
 class Rigify_Leg_FK2IK(bpy.types.Operator):
     """ Snaps an FK leg to an IK leg.
@@ -717,7 +748,11 @@ def execute(self, context):
         finally:
             context.user_preferences.edit.use_global_undo = use_global_undo
         return {'FINISHED'}
+''']
 
+REGISTER_OP_POLE = ['Rigify_Rot2PoleSwitch']
+
+UTILITIES_OP_POLE = ['''
 ###########################
 ## IK Rotation Pole Snap ##
 ###########################
@@ -743,7 +778,47 @@ def execute(self, context):
 
         rotPoleToggle(rig, self.limb_type, self.controls, self.ik_ctrl, self.fk_ctrl, self.parent, self.pole)
         return {'FINISHED'}
+''']
+
+REGISTER_RIG_ARM = REGISTER_OP_ARM_FKIK + REGISTER_OP_POLE
 
+UTILITIES_RIG_ARM = [
+    *UTILITIES_FUNC_ARM_FKIK,
+    *UTILITIES_FUNC_POLE,
+    *UTILITIES_OP_ARM_FKIK,
+    *UTILITIES_OP_POLE,
+]
+
+REGISTER_RIG_LEG = REGISTER_OP_LEG_FKIK + REGISTER_OP_POLE
+
+UTILITIES_RIG_LEG = [
+    *UTILITIES_FUNC_LEG_FKIK,
+    *UTILITIES_FUNC_POLE,
+    *UTILITIES_OP_LEG_FKIK,
+    *UTILITIES_OP_POLE,
+]
+
+##############################
+## Default set of utilities ##
+##############################
+
+UI_REGISTER = [
+    'RigUI',
+    'RigLayers',
+    *REGISTER_OP_ARM_FKIK,
+    *REGISTER_OP_LEG_FKIK,
+]
+
+# Include arm and leg utilities for now in case somebody wants to use
+# legacy limb rigs, which expect these to be available by default.
+UI_UTILITIES = [
+    *UTILITIES_FUNC_ARM_FKIK,
+    *UTILITIES_FUNC_LEG_FKIK,
+    *UTILITIES_OP_ARM_FKIK,
+    *UTILITIES_OP_LEG_FKIK,
+]
+
+UI_SLIDERS = '''
 ###################
 ## Rig UI Panels ##
 ###################
@@ -837,27 +912,3 @@ def draw(self, context):
     code += "        row.prop(context.active_object.data, 'layers', index=28, toggle=True, text='Root')\n"
 
     return code
-
-
-UI_REGISTER = '''
-
-def register():
-    bpy.utils.register_class(Rigify_Arm_FK2IK)
-    bpy.utils.register_class(Rigify_Arm_IK2FK)
-    bpy.utils.register_class(Rigify_Leg_FK2IK)
-    bpy.utils.register_class(Rigify_Leg_IK2FK)
-    bpy.utils.register_class(Rigify_Rot2PoleSwitch)
-    bpy.utils.register_class(RigUI)
-    bpy.utils.register_class(RigLayers)
-
-def unregister():
-    bpy.utils.unregister_class(Rigify_Arm_FK2IK)
-    bpy.utils.unregister_class(Rigify_Arm_IK2FK)
-    bpy.utils.unregister_class(Rigify_Leg_FK2IK)
-    bpy.utils.unregister_class(Rigify_Leg_IK2FK)
-    bpy.utils.register_class(Rigify_Rot2PoleSwitch)
-    bpy.utils.unregister_class(RigUI)
-    bpy.utils.unregister_class(RigLayers)
-
-register()
-'''

rigs/limbs/arm.py

@@ -9,6 +9,7 @@
 from ...utils       import MetarigError, make_mechanism_name, org
 from ...utils       import create_limb_widget, connected_children_names
 from ...utils       import align_bone_y_axis, align_bone_x_axis, align_bone_z_axis
+from ...rig_ui_template import UTILITIES_RIG_ARM, REGISTER_RIG_ARM
 from rna_prop_ui import rna_idprop_ui_prop_get
 from ..widgets import create_ikarrow_widget
 from math import trunc, pi
@@ -1067,7 +1068,11 @@ def generate(self):
         script += extra_script % (controls_string, bones['main_parent'], 'IK_follow',
                                   'pole_follow', 'pole_follow', 'root/parent', 'root/parent')
 
-        return [script]
+        return {
+            'script': [script],
+            'utilities': UTILITIES_RIG_ARM,
+            'register': REGISTER_RIG_ARM,
+        }
 
 
 def add_parameters(params):

rigs/limbs/leg.py

@@ -10,6 +10,7 @@
 from ...utils import MetarigError, make_mechanism_name, org
 from ...utils import create_limb_widget, connected_children_names
 from ...utils import align_bone_y_axis, align_bone_x_axis, align_bone_z_axis
+from ...rig_ui_template import UTILITIES_RIG_LEG, REGISTER_RIG_LEG
 from rna_prop_ui import rna_idprop_ui_prop_get
 from ..widgets import create_ikarrow_widget
 from math import trunc, pi
@@ -1390,7 +1391,11 @@ def generate(self):
         script += extra_script % (controls_string, bones['main_parent'], 'IK_follow',
                                   'pole_follow', 'pole_follow', 'root/parent', 'root/parent')
 
-        return [script]
+        return {
+            'script': [script],
+            'utilities': UTILITIES_RIG_LEG,
+            'register': REGISTER_RIG_LEG,
+        }
 
 
 def add_parameters(params):

rigs/limbs/paw.py

@@ -8,6 +8,7 @@
 from ...utils import MetarigError, make_mechanism_name, org
 from ...utils import create_limb_widget, connected_children_names
 from ...utils import align_bone_y_axis, align_bone_x_axis, align_bone_z_axis
+from ...rig_ui_template import UTILITIES_RIG_LEG, REGISTER_RIG_LEG
 from rna_prop_ui import rna_idprop_ui_prop_get
 from ..widgets import create_ikarrow_widget, create_gear_widget
 from ..widgets import create_foot_widget, create_ballsocket_widget
@@ -1218,7 +1219,11 @@ def generate(self):
         script += extra_script % (controls_string, bones['main_parent'], 'IK_follow',
                                   'pole_follow', 'pole_follow', 'root/parent', 'root/parent')
 
-        return [script]
+        return {
+            'script': [script],
+            'utilities': UTILITIES_RIG_LEG,
+            'register': REGISTER_RIG_LEG,
+        }
 
 
 def add_parameters(params):