Rigify: redesign generate.py and introduce a base rig class.

The main goals are to provide an official way for rigs to
interact in a structured way, and to remove mode switching
within rigs.

This involves introducing a base class for rigs that holds
rig-to-rig and rig-to-bone references, converting the main
generator into a class and passing it to rigs, and splitting
the single generate method into multiple passes.

For backward compatibility, old rigs are automatically handled
via a wrapper that translates between old and new API.

In addition, a way to create objects that receive the generate
callbacks that aren't rigs is introduced via the GeneratorPlugin
class. The UI script generation code is converted into a plugin.

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.

A new way to switch parents for IK bones based on the new
features is introduced, and used in the existing limb rigs.

Reviewers: icappiello campbellbarton

Differential Revision: https://developer.blender.org/D4624

1 files changed, 277 insertions, 55 deletions

diff --git a/rigify/utils/bones.py b/rigify/utils/bones.py
index 136ece7d..b5559a76 100644
--- a/rigify/utils/bones.py
+++ b/rigify/utils/bones.py
@@ -24,7 +24,8 @@ from mathutils import Vector, Matrix, Color
 from rna_prop_ui import rna_idprop_ui_prop_get
 
 from .errors import MetarigError
-from .naming import make_derived_name
+from .naming import get_name, make_derived_name
+from .misc import pairwise
 
 #=======================
 # Bone collection
@@ -55,7 +56,7 @@ class BoneDict(dict):
         raise ValueError("Invalid BoneDict value: %r" % (value))
 
     def __init__(self, *args, **kwargs):
-        super(BoneDict, self).__init__()
+        super().__init__()
 
         for key, value in dict(*args, **kwargs).items():
             dict.__setitem__(self, key, BoneDict.__sanitize_attr(key, value))
@@ -72,12 +73,14 @@ class BoneDict(dict):
         for key, value in dict(*args, **kwargs).items():
             dict.__setitem__(self, key, BoneDict.__sanitize_attr(key, value))
 
-    def flatten(self):
-        """Return all contained bones as a list."""
+    def flatten(self, key=None):
+        """Return all contained bones or a single key as a list."""
+
+        items = [self[key]] if key is not None else self.values()
 
         all_bones = []
 
-        for item in self.values():
+        for item in items:
             if isinstance(item, BoneDict):
                 all_bones.extend(item.flatten())
             elif isinstance(item, list):
@@ -90,6 +93,18 @@ class BoneDict(dict):
 #=======================
 # Bone manipulation
 #=======================
+#
+# NOTE: PREFER USING BoneUtilityMixin IN NEW STYLE RIGS!
+
+def get_bone(obj, bone_name):
+    """Get EditBone or PoseBone by name, depending on the current mode."""
+    if not bone_name:
+        return None
+    bones = obj.data.edit_bones if obj.mode == 'EDIT' else obj.pose.bones
+    if bone_name not in bones:
+        raise MetarigError("bone '%s' not found" % bone_name)
+    return bones[bone_name]
+
 
 def new_bone(obj, bone_name):
     """ Adds a new bone to the given armature object.
@@ -101,17 +116,14 @@ def new_bone(obj, bone_name):
         edit_bone.head = (0, 0, 0)
         edit_bone.tail = (0, 1, 0)
         edit_bone.roll = 0
-        bpy.ops.object.mode_set(mode='OBJECT')
-        bpy.ops.object.mode_set(mode='EDIT')
         return name
     else:
         raise MetarigError("Can't add new bone '%s' outside of edit mode" % bone_name)
 
 
-def copy_bone_simple(obj, bone_name, assign_name=''):
+def copy_bone(obj, bone_name, assign_name='', *, parent=False, bbone=False, length=None, scale=None):
     """ Makes a copy of the given bone in the given armature object.
-        but only copies head, tail positions and roll. Does not
-        address parenting either.
+        Returns the resulting bone's name.
     """
     #if bone_name not in obj.data.bones:
     if bone_name not in obj.data.edit_bones:
@@ -133,49 +145,36 @@ def copy_bone_simple(obj, bone_name, assign_name=''):
         edit_bone_2.tail = Vector(edit_bone_1.tail)
         edit_bone_2.roll = edit_bone_1.roll
 
-        return bone_name_2
-    else:
-        raise MetarigError("Cannot copy bones outside of edit mode")
-
-
-def copy_bone(obj, bone_name, assign_name=''):
-    """ Makes a copy of the given bone in the given armature object.
-        Returns the resulting bone's name.
-    """
-    #if bone_name not in obj.data.bones:
-    if bone_name not in obj.data.edit_bones:
-        raise MetarigError("copy_bone(): bone '%s' not found, cannot copy it" % bone_name)
+        if parent:
+            edit_bone_2.parent = edit_bone_1.parent
+            edit_bone_2.use_connect = edit_bone_1.use_connect
 
-    if obj == bpy.context.active_object and bpy.context.mode == 'EDIT_ARMATURE':
-        if assign_name == '':
-            assign_name = bone_name
-        # Copy the edit bone
-        edit_bone_1 = obj.data.edit_bones[bone_name]
-        edit_bone_2 = obj.data.edit_bones.new(assign_name)
-        bone_name_1 = bone_name
-        bone_name_2 = edit_bone_2.name
+            edit_bone_2.use_inherit_rotation = edit_bone_1.use_inherit_rotation
+            edit_bone_2.use_inherit_scale = edit_bone_1.use_inherit_scale
+            edit_bone_2.use_local_location = edit_bone_1.use_local_location
 
-        edit_bone_2.parent = edit_bone_1.parent
-        edit_bone_2.use_connect = edit_bone_1.use_connect
+        if bbone:
+            for name in ['bbone_segments',
+                         'bbone_easein', 'bbone_easeout',
+                         'bbone_rollin', 'bbone_rollout',
+                         'bbone_curveinx', 'bbone_curveiny', 'bbone_curveoutx', 'bbone_curveouty',
+                         'bbone_scaleinx', 'bbone_scaleiny', 'bbone_scaleoutx', 'bbone_scaleouty']:
+                setattr(edit_bone_2, name, getattr(edit_bone_1, name))
 
-        # Copy edit bone attributes
-        edit_bone_2.layers = list(edit_bone_1.layers)
+        # Resize the bone after copy if requested
+        if length is not None:
+            edit_bone_2.length = length
+        elif scale is not None:
+            edit_bone_2.length *= scale
 
-        edit_bone_2.head = Vector(edit_bone_1.head)
-        edit_bone_2.tail = Vector(edit_bone_1.tail)
-        edit_bone_2.roll = edit_bone_1.roll
-
-        edit_bone_2.use_inherit_rotation = edit_bone_1.use_inherit_rotation
-        edit_bone_2.use_inherit_scale = edit_bone_1.use_inherit_scale
-        edit_bone_2.use_local_location = edit_bone_1.use_local_location
-
-        edit_bone_2.use_deform = edit_bone_1.use_deform
-        edit_bone_2.bbone_segments = edit_bone_1.bbone_segments
-        edit_bone_2.bbone_easein = edit_bone_1.bbone_easein
-        edit_bone_2.bbone_easeout = edit_bone_1.bbone_easeout
+        return bone_name_2
+    else:
+        raise MetarigError("Cannot copy bones outside of edit mode")
 
-        bpy.ops.object.mode_set(mode='OBJECT')
 
+def copy_bone_properties(obj, bone_name_1, bone_name_2):
+    """ Copy transform and custom properties from bone 1 to bone 2. """
+    if obj.mode in {'OBJECT','POSE'}:
         # Get the pose bones
         pose_bone_1 = obj.pose.bones[bone_name_1]
         pose_bone_2 = obj.pose.bones[bone_name_2]
@@ -203,18 +202,24 @@ def copy_bone(obj, bone_name, assign_name=''):
                     prop2 = rna_idprop_ui_prop_get(pose_bone_2, key, create=True)
                     for key in prop1.keys():
                         prop2[key] = prop1[key]
+    else:
+        raise MetarigError("Cannot copy bone properties in edit mode")
 
-        bpy.ops.object.mode_set(mode='EDIT')
 
-        return bone_name_2
-    else:
-        raise MetarigError("Cannot copy bones outside of edit mode")
+def _legacy_copy_bone(obj, bone_name, assign_name=''):
+    """LEGACY ONLY, DON'T USE"""
+    new_name = copy_bone(obj, bone_name, assign_name, parent=True, bbone=True)
+    # Mode switch PER BONE CREATION?!
+    bpy.ops.object.mode_set(mode='OBJECT')
+    copy_bone_properties(obj, bone_name, new_name)
+    bpy.ops.object.mode_set(mode='EDIT')
+    return new_name
 
 
 def flip_bone(obj, bone_name):
     """ Flips an edit bone.
     """
-    if bone_name not in obj.data.bones:
+    if bone_name not in obj.data.edit_bones:
         raise MetarigError("flip_bone(): bone '%s' not found, cannot copy it" % bone_name)
 
     if obj == bpy.context.active_object and bpy.context.mode == 'EDIT_ARMATURE':
@@ -228,10 +233,38 @@ def flip_bone(obj, bone_name):
         raise MetarigError("Cannot flip bones outside of edit mode")
 
 
+def flip_bone_chain(obj, bone_names):
+    """Flips a connected bone chain."""
+    assert obj.mode == 'EDIT'
+
+    bones = [ obj.data.edit_bones[name] for name in bone_names ]
+
+    # Verify chain and unparent
+    for prev_bone, bone in pairwise(bones):
+        assert bone.parent == prev_bone and bone.use_connect
+
+    for bone in bones:
+        bone.parent = None
+        bone.use_connect = False
+        for child in bone.children:
+            child.use_connect = False
+
+    # Flip bones
+    for bone in bones:
+        head, tail = Vector(bone.head), Vector(bone.tail)
+        bone.tail = head + tail
+        bone.head, bone.tail = tail, head
+
+    # Re-parent
+    for bone, next_bone in pairwise(bones):
+        bone.parent = next_bone
+        bone.use_connect = True
+
+
 def put_bone(obj, bone_name, pos):
     """ Places a bone at the given position.
     """
-    if bone_name not in obj.data.bones:
+    if bone_name not in obj.data.edit_bones:
         raise MetarigError("put_bone(): bone '%s' not found, cannot move it" % bone_name)
 
     if obj == bpy.context.active_object and bpy.context.mode == 'EDIT_ARMATURE':
@@ -243,7 +276,14 @@ def put_bone(obj, bone_name, pos):
         raise MetarigError("Cannot 'put' bones outside of edit mode")
 
 
-def make_nonscaling_child(obj, bone_name, location, child_name_postfix=""):
+def disable_bbones(obj, bone_names):
+    """Disables B-Bone segments on the specified bones."""
+    assert(obj.mode != 'EDIT')
+    for bone in bone_names:
+        obj.data.bones[bone].bbone_segments = 1
+
+
+def _legacy_make_nonscaling_child(obj, bone_name, location, child_name_postfix=""):
     """ Takes the named bone and creates a non-scaling child of it at
         the given location.  The returned bone (returned by name) is not
         a true child, but behaves like one sans inheriting scaling.
@@ -251,8 +291,10 @@ def make_nonscaling_child(obj, bone_name, location, child_name_postfix=""):
         It is intended as an intermediate construction to prevent rig types
         from scaling with their parents.  The named bone is assumed to be
         an ORG bone.
+
+        LEGACY ONLY, DON'T USE
     """
-    if bone_name not in obj.data.bones:
+    if bone_name not in obj.data.edit_bones:
         raise MetarigError("make_nonscaling_child(): bone '%s' not found, cannot copy it" % bone_name)
 
     if obj == bpy.context.active_object and bpy.context.mode == 'EDIT_ARMATURE':
@@ -305,11 +347,129 @@ def make_nonscaling_child(obj, bone_name, location, child_name_postfix=""):
         raise MetarigError("Cannot make nonscaling child outside of edit mode")
 
 
+#===================================
+# Bone manipulation as rig methods
+#===================================
+
+
+class BoneUtilityMixin(object):
+    """
+    Provides methods for more convenient creation of bones.
+
+    Requires self.obj to be the armature object being worked on.
+    """
+    def register_new_bone(self, new_name, old_name=None):
+        """Registers creation or renaming of a bone based on old_name"""
+        pass
+
+    def new_bone(self, new_name):
+        """Create a new bone with the specified name."""
+        name = new_bone(self.obj, bone_name)
+        self.register_new_bone(self, name)
+        return name
+
+    def copy_bone(self, bone_name, new_name='', *, parent=False, bbone=False, length=None, scale=None):
+        """Copy the bone with the given name, returning the new name."""
+        name = copy_bone(self.obj, bone_name, new_name, parent=parent, bbone=bbone, length=length, scale=scale)
+        self.register_new_bone(name, bone_name)
+        return name
+
+    def copy_bone_properties(self, src_name, tgt_name):
+        """Copy pose-mode properties of the bone."""
+        copy_bone_properties(self.obj, src_name, tgt_name)
+
+    def rename_bone(self, old_name, new_name):
+        """Rename the bone, returning the actual new name."""
+        bone = self.get_bone(old_name)
+        bone.name = new_name
+        if bone.name != old_name:
+            self.register_new_bone(bone.name, old_name)
+        return bone.name
+
+    def get_bone(self, bone_name):
+        """Get EditBone or PoseBone by name, depending on the current mode."""
+        return get_bone(self.obj, bone_name)
+
+    def get_bone_parent(self, bone_name):
+        """Get the name of the parent bone, or None."""
+        return get_name(self.get_bone(bone_name).parent)
+
+    def set_bone_parent(self, bone_name, parent_name, use_connect=False):
+        """Set the parent of the bone."""
+        eb = self.obj.data.edit_bones
+        bone = eb[bone_name]
+        if use_connect is not None:
+            bone.use_connect = use_connect
+        bone.parent = (eb[parent_name] if parent_name else None)
+
+    def parent_bone_chain(self, bone_names, use_connect=None):
+        """Link bones into a chain with parenting. First bone may be None."""
+        for parent, child in pairwise(bone_names):
+            self.set_bone_parent(child, parent, use_connect=use_connect)
+
+#=============================================
+# B-Bones
+#=============================================
+
+def connect_bbone_chain_handles(obj, bone_names):
+    assert obj.mode == 'EDIT'
+
+    for prev_name, next_name in pairwise(bone_names):
+        prev_bone = get_bone(obj, prev_name)
+        next_bone = get_bone(obj, next_name)
+
+        prev_bone.bbone_handle_type_end = 'ABSOLUTE'
+        prev_bone.bbone_custom_handle_end = next_bone
+
+        next_bone.bbone_handle_type_start = 'ABSOLUTE'
+        next_bone.bbone_custom_handle_start = prev_bone
+
 #=============================================
 # Math
 #=============================================
 
 
+def is_same_position(obj, bone_name1, bone_name2):
+    head1 = get_bone(obj, bone_name1).head
+    head2 = get_bone(obj, bone_name2).head
+
+    return (head1 - head2).length < 1e-5
+
+
+def is_connected_position(obj, bone_name1, bone_name2):
+    tail1 = get_bone(obj, bone_name1).tail
+    head2 = get_bone(obj, bone_name2).head
+
+    return (tail1 - head2).length < 1e-5
+
+
+def copy_bone_position(obj, bone_name, target_bone_name, *, length=None, scale=None):
+    """ Completely copies the position and orientation of the bone. """
+    bone1_e = obj.data.edit_bones[bone_name]
+    bone2_e = obj.data.edit_bones[target_bone_name]
+
+    bone2_e.head = bone1_e.head
+    bone2_e.tail = bone1_e.tail
+    bone2_e.roll = bone1_e.roll
+
+    # Resize the bone after copy if requested
+    if length is not None:
+        bone2_e.length = length
+    elif scale is not None:
+        bone2_e.length *= scale
+
+
+def align_bone_orientation(obj, bone_name, target_bone_name):
+    """ Aligns the orientation of bone to target bone. """
+    bone1_e = obj.data.edit_bones[bone_name]
+    bone2_e = obj.data.edit_bones[target_bone_name]
+
+    axis = bone2_e.y_axis.normalized() * bone1_e.length
+
+    bone1_e.tail = bone1_e.head + axis
+    bone1_e.roll = bone2_e.roll
+
+
 def align_bone_roll(obj, bone1, bone2):
     """ Aligns the roll of two bones.
     """
@@ -416,3 +576,65 @@ def align_bone_y_axis(obj, bone, vec):
     vec = vec * bone_e.length
 
     bone_e.tail = bone_e.head + vec
+
+
+def compute_chain_x_axis(obj, bone_names):
+    """
+    Compute the x axis of all bones to be perpendicular
+    to the primary plane in which the bones lie.
+    """
+    eb = obj.data.edit_bones
+
+    assert(len(bone_names) > 1)
+    first_bone = eb[bone_names[0]]
+    last_bone = eb[bone_names[-1]]
+
+    # Compute normal to the plane defined by the first bone,
+    # and the end of the last bone in the chain
+    chain_y_axis = last_bone.tail - first_bone.head
+    chain_rot_axis = first_bone.y_axis.cross(chain_y_axis)
+
+    if chain_rot_axis.length < first_bone.length/100:
+        return first_bone.x_axis.normalized()
+    else:
+        return chain_rot_axis.normalized()
+
+
+def align_chain_x_axis(obj, bone_names):
+    """
+    Aligns the x axis of all bones to be perpendicular
+    to the primary plane in which the bones lie.
+    """
+    chain_rot_axis = compute_chain_x_axis(obj, bone_names)
+
+    for name in bone_names:
+        align_bone_x_axis(obj, name, chain_rot_axis)
+
+
+def align_bone_to_axis(obj, bone_name, axis, *, length=None, roll=0, flip=False):
+    """
+    Aligns the Y axis of the bone to the global axis (x,y,z,-x,-y,-z),
+    optionally adjusting length and initially flipping the bone.
+    """
+    bone_e = obj.data.edit_bones[bone_name]
+
+    if length is None:
+        length = bone_e.length
+    if roll is None:
+        roll = bone_e.roll
+
+    if axis[0] == '-':
+        length = -length
+        axis = axis[1:]
+
+    vec = Vector((0,0,0))
+    setattr(vec, axis, length)
+
+    if flip:
+        base = Vector(bone_e.tail)
+        bone_e.tail = base + vec
+        bone_e.head = base
+    else:
+        bone_e.tail = bone_e.head + vec
+
+    bone_e.roll = roll