Rigify:

- Added IK/FK snapping (both directions) for legs.
- Cleaned up another operator so that it works with undo.
- PEP8 cleanups. 

7 files changed, 254 insertions, 48 deletions

diff --git a/rigify/__init__.py b/rigify/__init__.py
index a235b08c..a20fd58d 100644
--- a/rigify/__init__.py
+++ b/rigify/__init__.py
@@ -116,7 +116,6 @@ class RigifyParameters(bpy.types.PropertyGroup):
     name = bpy.props.StringProperty()
 
 
-
 ##### REGISTER #####
 
 def register():
@@ -142,6 +141,7 @@ def register():
         except AttributeError:
             pass
 
+
 def unregister():
     del bpy.types.PoseBone.rigify_type
     del bpy.types.PoseBone.rigify_parameters
diff --git a/rigify/rig_ui_template.py b/rigify/rig_ui_template.py
index ac0e919b..c9f4bdfc 100644
--- a/rigify/rig_ui_template.py
+++ b/rigify/rig_ui_template.py
@@ -31,15 +31,15 @@ def get_pose_matrix_in_other_space(mat, pbs):
         This is with constraints applied.
     """
     rest_inv = pbs.bone.matrix_local.inverted()
-    
+
     if pbs.parent:
         par_inv = pbs.parent.matrix.inverted()
         par_rest = pbs.parent.bone.matrix_local.copy()
-        
+
         smat = rest_inv * (par_rest * (par_inv * mat))
     else:
         smat = rest_inv * mat
-    
+
     return smat
 
 
@@ -55,7 +55,7 @@ def set_pose_rotation(pb, mat):
         Matrix should be given in local space.
     """
     q = mat.to_quaternion()
-    
+
     if pb.rotation_mode == 'QUATERNION':
         pb.rotation_quaternion = q
     elif pb.rotation_mode == 'AXIS_ANGLE':
@@ -82,63 +82,63 @@ def set_pose_translation(pb, mat):
         pb.location = mat.to_translation()
     else:
         loc = mat.to_translation()
-        
+
         rest = pb.bone.matrix_local.copy()
         if pb.bone.parent:
             par_rest = pb.bone.parent.matrix_local.copy()
         else:
             par_rest = Matrix()
-        
+
         q = (par_rest.inverted() * rest).to_quaternion()
         pb.location = loc * q
 
 
-def fk2ik(obj, fk, ik):
+def fk2ik_arm(obj, fk, ik):
     """ Matches the fk bones in an arm rig to the ik bones.
     """
     pb = obj.pose.bones
-    
+
     uarm = pb[fk[0]]
     farm = pb[fk[1]]
     hand = pb[fk[2]]
-    
+
     uarmi = pb[ik[0]]
     farmi = pb[ik[1]]
     handi = pb[ik[2]]
-    
+
     uarmmat = get_pose_matrix_in_other_space(uarmi.matrix, uarm)
     set_pose_rotation(uarm, uarmmat)
     set_pose_scale(uarm, uarmmat)
     bpy.ops.object.mode_set(mode='OBJECT')
     bpy.ops.object.mode_set(mode='POSE')
-    
+
     farmmat = get_pose_matrix_in_other_space(farmi.matrix, farm)
     set_pose_rotation(farm, farmmat)
     set_pose_scale(farm, farmmat)
     bpy.ops.object.mode_set(mode='OBJECT')
     bpy.ops.object.mode_set(mode='POSE')
-    
+
     handmat = get_pose_matrix_in_other_space(handi.matrix, hand)
     set_pose_rotation(hand, handmat)
     set_pose_scale(hand, handmat)
     bpy.ops.object.mode_set(mode='OBJECT')
     bpy.ops.object.mode_set(mode='POSE')
-    
-    
-def ik2fk(obj, fk, ik):
+
+
+def ik2fk_arm(obj, fk, ik):
     """ Matches the ik bones in an arm rig to the fk bones.
     """
     pb = obj.pose.bones
-    
+
     uarm = pb[fk[0]]
     farm = pb[fk[1]]
     hand = pb[fk[2]]
-    
+
     uarmi = pb[ik[0]]
     farmi = pb[ik[1]]
     handi = pb[ik[2]]
     pole = pb[ik[3]]
-    
+
     # Hand position
     handmat = get_pose_matrix_in_other_space(hand.matrix, handi)
     set_pose_translation(handi, handmat)
@@ -146,15 +146,15 @@ def ik2fk(obj, fk, ik):
     set_pose_scale(handi, handmat)
     bpy.ops.object.mode_set(mode='OBJECT')
     bpy.ops.object.mode_set(mode='POSE')
-    
+
     # Pole target position
     a = uarm.matrix.to_translation()
     b = farm.matrix.to_translation() + farm.vector
-    
+
     # Vector from the head of the upper arm to the
     # tip of the forearm
     armv = b - a
-    
+
     # Create a vector that is not aligned with armv.
     # It doesn't matter what vector.  Just any vector
     # that's guaranteed to not be pointing in the same
@@ -165,33 +165,33 @@ def ik2fk(obj, fk, ik):
         v = Vector((0,1,0))
     else:
         v = Vector((0,0,1))
-    
+
     # cross v with armv to get a vector perpendicular to armv
     pv = v.cross(armv).normalized() * (uarm.length + farm.length)
-    
+
     def set_pole(pvi):
         # Translate pvi into armature space
         ploc = a + (armv/2) + pvi
-        
+
         # Set pole target to location
         mat = get_pose_matrix_in_other_space(Matrix().Translation(ploc), pole)
         set_pose_translation(pole, mat)
-        
+
         bpy.ops.object.mode_set(mode='OBJECT')
         bpy.ops.object.mode_set(mode='POSE')
-    
+
     set_pole(pv)
-    
+
     # Get the rotation difference between the ik and fk upper arms
     q1 = uarm.matrix.to_quaternion()
     q2 = uarmi.matrix.to_quaternion()
     angle = acos(min(1,max(-1,q1.dot(q2)))) * 2
-    
+
     # Compensate for the rotation difference
     if angle > 0.00001:
         pv *= Matrix.Rotation(angle, 4, armv).to_quaternion()
         set_pole(pv)
-    
+
         q1 = uarm.matrix.to_quaternion()
         q2 = uarmi.matrix.to_quaternion()
         angle2 = acos(min(1,max(-1,q1.dot(q2)))) * 2
@@ -200,17 +200,132 @@ def ik2fk(obj, fk, ik):
             set_pole(pv)
 
 
+def fk2ik_leg(obj, fk, ik):
+    """ Matches the fk bones in a leg rig to the ik bones.
+    """
+    pb = obj.pose.bones
+
+    thigh = pb[fk[0]]
+    shin  = pb[fk[1]]
+    foot  = pb[fk[2]]
+    mfoot = pb[fk[3]]
+
+    thighi = pb[ik[0]]
+    shini  = pb[ik[1]]
+    mfooti = pb[ik[2]]
+
+    thighmat = get_pose_matrix_in_other_space(thighi.matrix, thigh)
+    set_pose_rotation(thigh, thighmat)
+    set_pose_scale(thigh, thighmat)
+    bpy.ops.object.mode_set(mode='OBJECT')
+    bpy.ops.object.mode_set(mode='POSE')
+
+    shinmat = get_pose_matrix_in_other_space(shini.matrix, shin)
+    set_pose_rotation(shin, shinmat)
+    set_pose_scale(shin, shinmat)
+    bpy.ops.object.mode_set(mode='OBJECT')
+    bpy.ops.object.mode_set(mode='POSE')
+
+    mat = mfoot.bone.matrix_local.inverted() * foot.bone.matrix_local
+    footmat = get_pose_matrix_in_other_space(mfooti.matrix, foot) * mat
+    set_pose_rotation(foot, footmat)
+    set_pose_scale(foot, footmat)
+    bpy.ops.object.mode_set(mode='OBJECT')
+    bpy.ops.object.mode_set(mode='POSE')
+
+
+def ik2fk_leg(obj, fk, ik):
+    """ Matches the ik bones in a leg rig to the fk bones.
+    """
+    pb = obj.pose.bones
+
+    thigh = pb[fk[0]]
+    shin  = pb[fk[1]]
+    mfoot = pb[fk[2]]
+
+    thighi   = pb[ik[0]]
+    shini    = pb[ik[1]]
+    footi    = pb[ik[2]]
+    footroll = pb[ik[3]]
+    pole     = pb[ik[4]]
+    mfooti   = pb[ik[5]]
+
+    # Clear footroll
+    set_pose_rotation(footroll, Matrix())
+
+    # Foot position
+    mat = mfooti.bone.matrix_local.inverted() * footi.bone.matrix_local
+    footmat = get_pose_matrix_in_other_space(mfoot.matrix, footi) * mat
+    set_pose_translation(footi, footmat)
+    set_pose_rotation(footi, footmat)
+    set_pose_scale(footi, footmat)
+    bpy.ops.object.mode_set(mode='OBJECT')
+    bpy.ops.object.mode_set(mode='POSE')
+
+    # Pole target position
+    a = thigh.matrix.to_translation()
+    b = shin.matrix.to_translation() + shin.vector
+
+    # Vector from the head of the thigh to the
+    # tip of the shin
+    legv = b - a
+
+    # Create a vector that is not aligned with legv.
+    # It doesn't matter what vector.  Just any vector
+    # that's guaranteed to not be pointing in the same
+    # direction.
+    if abs(legv[0]) < abs(legv[1]) and abs(legv[0]) < abs(legv[2]):
+        v = Vector((1,0,0))
+    elif abs(legv[1]) < abs(legv[2]):
+        v = Vector((0,1,0))
+    else:
+        v = Vector((0,0,1))
+
+    # Get a vector perpendicular to legv
+    pv = v.cross(legv).normalized() * (thigh.length + shin.length)
+
+    def set_pole(pvi):
+        # Translate pvi into armature space
+        ploc = a + (legv/2) + pvi
+
+        # Set pole target to location
+        mat = get_pose_matrix_in_other_space(Matrix().Translation(ploc), pole)
+        set_pose_translation(pole, mat)
+
+        bpy.ops.object.mode_set(mode='OBJECT')
+        bpy.ops.object.mode_set(mode='POSE')
+
+    set_pole(pv)
+
+    # Get the rotation difference between the ik and fk thighs
+    q1 = thigh.matrix.to_quaternion()
+    q2 = thighi.matrix.to_quaternion()
+    angle = acos(min(1,max(-1,q1.dot(q2)))) * 2
+
+    # Compensate for the rotation difference
+    if angle > 0.00001:
+        pv *= Matrix.Rotation(angle, 4, legv).to_quaternion()
+        set_pole(pv)
+
+        q1 = thigh.matrix.to_quaternion()
+        q2 = thighi.matrix.to_quaternion()
+        angle2 = acos(min(1,max(-1,q1.dot(q2)))) * 2
+        if angle2 > 0.00001:
+            pv *= Matrix.Rotation((angle*(-2)), 4, legv).to_quaternion()
+            set_pole(pv)
+
+
 class Rigify_Arm_FK2IK(bpy.types.Operator):
     """ Snaps an FK arm to an IK arm.
     """
     bl_idname = "pose.rigify_arm_fk2ik_" + rig_id
     bl_label = "Rigify Snap FK arm to IK"
     bl_options = {'UNDO'}
-    
+
     uarm_fk = bpy.props.StringProperty(name="Upper Arm FK Name")
     farm_fk = bpy.props.StringProperty(name="Forerm FK Name")
     hand_fk = bpy.props.StringProperty(name="Hand FK Name")
-    
+
     uarm_ik = bpy.props.StringProperty(name="Upper Arm IK Name")
     farm_ik = bpy.props.StringProperty(name="Forearm IK Name")
     hand_ik = bpy.props.StringProperty(name="Hand IK Name")
@@ -223,7 +338,7 @@ class Rigify_Arm_FK2IK(bpy.types.Operator):
         use_global_undo = context.user_preferences.edit.use_global_undo
         context.user_preferences.edit.use_global_undo = False
         try:
-            fk2ik(context.active_object, fk=[self.uarm_fk, self.farm_fk, self.hand_fk], ik=[self.uarm_ik, self.farm_ik, self.hand_ik])
+            fk2ik_arm(context.active_object, fk=[self.uarm_fk, self.farm_fk, self.hand_fk], ik=[self.uarm_ik, self.farm_ik, self.hand_ik])
         finally:
             context.user_preferences.edit.use_global_undo = use_global_undo
         return {'FINISHED'}
@@ -233,13 +348,13 @@ class Rigify_Arm_IK2FK(bpy.types.Operator):
     """ Snaps an IK arm to an FK arm.
     """
     bl_idname = "pose.rigify_arm_ik2fk_" + rig_id
-    bl_label = "Snap IK arm to FK"
+    bl_label = "Rigify Snap IK arm to FK"
     bl_options = {'UNDO'}
-    
+
     uarm_fk = bpy.props.StringProperty(name="Upper Arm FK Name")
     farm_fk = bpy.props.StringProperty(name="Forerm FK Name")
     hand_fk = bpy.props.StringProperty(name="Hand FK Name")
-    
+
     uarm_ik = bpy.props.StringProperty(name="Upper Arm IK Name")
     farm_ik = bpy.props.StringProperty(name="Forearm IK Name")
     hand_ik = bpy.props.StringProperty(name="Hand IK Name")
@@ -253,7 +368,70 @@ class Rigify_Arm_IK2FK(bpy.types.Operator):
         use_global_undo = context.user_preferences.edit.use_global_undo
         context.user_preferences.edit.use_global_undo = False
         try:
-            ik2fk(context.active_object, fk=[self.uarm_fk, self.farm_fk, self.hand_fk], ik=[self.uarm_ik, self.farm_ik, self.hand_ik, self.pole])
+            ik2fk_arm(context.active_object, fk=[self.uarm_fk, self.farm_fk, self.hand_fk], ik=[self.uarm_ik, self.farm_ik, self.hand_ik, self.pole])
+        finally:
+            context.user_preferences.edit.use_global_undo = use_global_undo
+        return {'FINISHED'}
+
+
+class Rigify_Leg_FK2IK(bpy.types.Operator):
+    """ Snaps an FK leg to an IK leg.
+    """
+    bl_idname = "pose.rigify_leg_fk2ik_" + rig_id
+    bl_label = "Rigify Snap FK leg to IK"
+    bl_options = {'UNDO'}
+
+    thigh_fk = bpy.props.StringProperty(name="Thigh FK Name")
+    shin_fk  = bpy.props.StringProperty(name="Shin FK Name")
+    foot_fk  = bpy.props.StringProperty(name="Foot FK Name")
+    mfoot_fk = bpy.props.StringProperty(name="MFoot FK Name")
+
+
+    thigh_ik = bpy.props.StringProperty(name="Thigh IK Name")
+    shin_ik  = bpy.props.StringProperty(name="Shin IK Name")
+    mfoot_ik = bpy.props.StringProperty(name="MFoot IK Name")
+
+    @classmethod
+    def poll(cls, context):
+        return (context.active_object != None and context.mode == 'POSE')
+
+    def execute(self, context):
+        use_global_undo = context.user_preferences.edit.use_global_undo
+        context.user_preferences.edit.use_global_undo = False
+        try:
+            fk2ik_leg(context.active_object, fk=[self.thigh_fk, self.shin_fk, self.foot_fk, self.mfoot_fk], ik=[self.thigh_ik, self.shin_ik, self.mfoot_ik])
+        finally:
+            context.user_preferences.edit.use_global_undo = use_global_undo
+        return {'FINISHED'}
+
+
+class Rigify_Leg_IK2FK(bpy.types.Operator):
+    """ Snaps an IK leg to an FK leg.
+    """
+    bl_idname = "pose.rigify_leg_ik2fk_" + rig_id
+    bl_label = "Rigify Snap IK leg to FK"
+    bl_options = {'UNDO'}
+
+    thigh_fk = bpy.props.StringProperty(name="Thigh FK Name")
+    shin_fk  = bpy.props.StringProperty(name="Shin FK Name")
+    mfoot_fk = bpy.props.StringProperty(name="MFoot FK Name")
+
+    thigh_ik = bpy.props.StringProperty(name="Thigh IK Name")
+    shin_ik  = bpy.props.StringProperty(name="Shin IK Name")
+    foot_ik  = bpy.props.StringProperty(name="Foot IK Name")
+    footroll = bpy.props.StringProperty(name="Foot Roll Name")
+    pole     = bpy.props.StringProperty(name="Pole IK Name")
+    mfoot_ik = bpy.props.StringProperty(name="MFoot IK Name")
+
+    @classmethod
+    def poll(cls, context):
+        return (context.active_object != None and context.mode == 'POSE')
+
+    def execute(self, context):
+        use_global_undo = context.user_preferences.edit.use_global_undo
+        context.user_preferences.edit.use_global_undo = False
+        try:
+            ik2fk_leg(context.active_object, fk=[self.thigh_fk, self.shin_fk, self.mfoot_fk], ik=[self.thigh_ik, self.shin_ik, self.foot_ik, self.footroll, self.pole, self.mfoot_ik])
         finally:
             context.user_preferences.edit.use_global_undo = use_global_undo
         return {'FINISHED'}
@@ -261,6 +439,8 @@ class Rigify_Arm_IK2FK(bpy.types.Operator):
 
 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)
 
 
 class RigUI(bpy.types.Panel):
diff --git a/rigify/rigs/biped/leg/__init__.py b/rigify/rigs/biped/leg/__init__.py
index 65af500a..f2ea2e51 100644
--- a/rigify/rigs/biped/leg/__init__.py
+++ b/rigify/rigs/biped/leg/__init__.py
@@ -26,12 +26,35 @@ imp.reload(ik)
 imp.reload(deform)
 
 script = """
-fk_leg = ["%s", "%s", "%s"]
-ik_leg = ["%s", "%s", "%s"]
+fk_leg = ["%s", "%s", "%s", "%s"]
+ik_leg = ["%s", "%s", "%s", "%s", "%s", "%s"]
 if is_selected(fk_leg+ik_leg):
-    layout.prop(pose_bones[ik_leg[0]], '["ikfk_switch"]', text="FK / IK (" + ik_leg[0] + ")", slider=True)
+    layout.prop(pose_bones[ik_leg[2]], '["ikfk_switch"]', text="FK / IK (" + ik_leg[2] + ")", slider=True)
 if is_selected(fk_leg):
-    layout.prop(pose_bones[fk_leg[0]], '["isolate"]', text="Isolate Rotation (" + fk_leg[0] + ")", slider=True)
+    try:
+        pose_bones[fk_leg[0]]["isolate"]
+        layout.prop(pose_bones[fk_leg[0]], '["isolate"]', text="Isolate Rotation (" + fk_leg[0] + ")", slider=True)
+    except KeyError:
+        pass
+if is_selected(fk_leg+ik_leg):
+    p = layout.operator("pose.rigify_leg_fk2ik_" + rig_id, text="Snap FK->IK (" + fk_leg[0] + ")")
+    p.thigh_fk = fk_leg[0]
+    p.shin_fk  = fk_leg[1]
+    p.foot_fk  = fk_leg[2]
+    p.mfoot_fk = fk_leg[3]
+    p.thigh_ik = ik_leg[0]
+    p.shin_ik  = ik_leg[1]
+    p.mfoot_ik = ik_leg[5]
+    p = layout.operator("pose.rigify_leg_ik2fk_" + rig_id, text="Snap IK->FK (" + fk_leg[0] + ")")
+    p.thigh_fk  = fk_leg[0]
+    p.shin_fk   = fk_leg[1]
+    p.mfoot_fk  = fk_leg[3]
+    p.thigh_ik  = ik_leg[0]
+    p.shin_ik   = ik_leg[1]
+    p.foot_ik   = ik_leg[2]
+    p.pole      = ik_leg[3]
+    p.footroll  = ik_leg[4]
+    p.mfoot_ik  = ik_leg[5]
 """
 
 
@@ -64,7 +87,7 @@ class Rig:
         self.deform_rig.generate()
         fk_controls = self.fk_rig.generate()
         ik_controls = self.ik_rig.generate()
-        return [script % (fk_controls[0], fk_controls[1], fk_controls[2], ik_controls[0], ik_controls[1], ik_controls[2])]
+        return [script % (fk_controls[0], fk_controls[1], fk_controls[2], fk_controls[3], ik_controls[0], ik_controls[1], ik_controls[2], ik_controls[3], ik_controls[4], ik_controls[5])]
 
     @classmethod
     def add_parameters(self, group):
diff --git a/rigify/rigs/biped/leg/fk.py b/rigify/rigs/biped/leg/fk.py
index f32fc36d..bc7a8dbd 100644
--- a/rigify/rigs/biped/leg/fk.py
+++ b/rigify/rigs/biped/leg/fk.py
@@ -242,5 +242,5 @@ class Rig:
             mod = ob.modifiers.new("subsurf", 'SUBSURF')
             mod.levels = 2
 
-        return [thigh, shin, foot]
+        return [thigh, shin, foot, foot_mch]
 
diff --git a/rigify/rigs/biped/leg/ik.py b/rigify/rigs/biped/leg/ik.py
index b2328f4d..42ead8b5 100644
--- a/rigify/rigs/biped/leg/ik.py
+++ b/rigify/rigs/biped/leg/ik.py
@@ -111,7 +111,6 @@ class Rig:
                     heel = b.name
                     if len(b.children) > 0:
                         rocker = b.children[0].name
-                    
 
         if foot == None or heel == None:
             print("blah")
@@ -300,7 +299,6 @@ class Rig:
                 flip_bone(self.obj, rocker2)
             else:
                 flip_bone(self.obj, rocker1)
-            
 
         # Weird alignment issues.  Fix.
         toe_parent_e.head = Vector(org_foot_e.head)
@@ -574,5 +572,5 @@ class Rig:
             mod = ob.modifiers.new("subsurf", 'SUBSURF')
             mod.levels = 2
 
-        return [foot, pole, foot_roll]
+        return [thigh, shin, foot, pole, foot_roll, foot_ik_target]
 
diff --git a/rigify/rigs/finger.py b/rigify/rigs/finger.py
index 0676899b..0bcea44b 100644
--- a/rigify/rigs/finger.py
+++ b/rigify/rigs/finger.py
@@ -42,7 +42,6 @@ class Rig:
         if len(self.org_bones) <= 1:
             raise MetarigError("RIGIFY ERROR: Bone '%s': input to rig type must be a chain of 2 or more bones." % (strip_org(bone)))
 
-
         # Get user-specified layers, if they exist
         if params.separate_extra_layers:
             self.ex_layers = list(params.extra_layers)
diff --git a/rigify/ui.py b/rigify/ui.py
index c9f1e472..dd91c8cc 100644
--- a/rigify/ui.py
+++ b/rigify/ui.py
@@ -216,11 +216,15 @@ class Sample(bpy.types.Operator):
 
     bl_idname = "armature.metarig_sample_add"
     bl_label = "Add a sample metarig for a rig type"
+    bl_options = {'UNDO'}
 
     metarig_type = StringProperty(name="Type", description="Name of the rig type to generate a sample of", maxlen=128, default="")
 
+
     def execute(self, context):
         if context.mode == 'EDIT_ARMATURE' and self.metarig_type != "":
+            use_global_undo = context.user_preferences.edit.use_global_undo
+            context.user_preferences.edit.use_global_undo = False
             try:
                 rig = get_rig_type(self.metarig_type).Rig
                 create_sample = rig.create_sample
@@ -228,7 +232,9 @@ class Sample(bpy.types.Operator):
                 print("Rigify: rig type has no sample.")
             else:
                 create_sample(context.active_object)
-            bpy.ops.object.mode_set(mode='EDIT')
+            finally:
+                context.user_preferences.edit.use_global_undo = use_global_undo
+                bpy.ops.object.mode_set(mode='EDIT')
 
         return {'FINISHED'}