Rigify: implement optional custom pivot controls.

- Add an optional custom pivot between torso and the rest of the spine.
- Add a custom pivot rig that can be used as a parent of the spine.
- Add an optional custom pivot under limb IK controls.

6 files changed, 241 insertions, 22 deletions

diff --git a/rigify/rigs/basic/pivot.py b/rigify/rigs/basic/pivot.py
new file mode 100644
index 00000000..d4c5898d
--- /dev/null
+++ b/rigify/rigs/basic/pivot.py
@@ -0,0 +1,143 @@
+#====================== BEGIN GPL LICENSE BLOCK ======================
+#
+#  This program is free software; you can redistribute it and/or
+#  modify it under the terms of the GNU General Public License
+#  as published by the Free Software Foundation; either version 2
+#  of the License, or (at your option) any later version.
+#
+#  This program is distributed in the hope that it will be useful,
+#  but WITHOUT ANY WARRANTY; without even the implied warranty of
+#  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+#  GNU General Public License for more details.
+#
+#  You should have received a copy of the GNU General Public License
+#  along with this program; if not, write to the Free Software Foundation,
+#  Inc., 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301, USA.
+#
+#======================= END GPL LICENSE BLOCK ========================
+
+# <pep8 compliant>
+
+import bpy
+
+from ...base_rig import BaseRig
+
+from ...utils.naming import make_derived_name
+from ...utils.widgets_basic import create_cube_widget, create_pivot_widget
+
+
+class Rig(BaseRig):
+    """ A rig providing a rotation pivot control that can be moved. """
+    def find_org_bones(self, pose_bone):
+        return pose_bone.name
+
+
+    def initialize(self):
+        self.make_control = self.params.make_extra_control
+        self.make_deform  = self.params.make_extra_deform
+
+
+    def generate_bones(self):
+        org = self.bones.org
+
+        self.bones.ctrl.pivot = self.copy_bone(org, make_derived_name(org, 'ctrl'), parent=not self.make_control)
+
+        if self.make_control:
+            self.bones.ctrl.master = self.copy_bone(org, make_derived_name(org, 'ctrl', '_master'), parent=True)
+
+        if self.make_deform:
+            self.bones.deform = self.copy_bone(org, make_derived_name(org, 'def'), bbone=True)
+
+
+    def parent_bones(self):
+        if self.make_control:
+            self.set_bone_parent(self.bones.ctrl.pivot, self.bones.ctrl.master, use_connect=False)
+
+        self.set_bone_parent(self.bones.org, self.bones.ctrl.pivot, use_connect=False)
+
+        if self.make_deform:
+            self.set_bone_parent(self.bones.deform, self.bones.org, use_connect=False)
+
+
+    def configure_bones(self):
+        self.copy_bone_properties(self.bones.org, self.bones.ctrl.pivot)
+
+        if self.make_control:
+            self.copy_bone_properties(self.bones.org, self.bones.ctrl.master)
+
+
+    def rig_bones(self):
+        self.make_constraint(
+            self.bones.org, 'COPY_LOCATION', self.bones.ctrl.pivot,
+            space='LOCAL', invert_xyz=(True,)*3
+        )
+
+
+    def generate_widgets(self):
+        create_pivot_widget(self.obj, self.bones.ctrl.pivot, square=True, axis_size=2.0)
+
+        if self.make_control:
+            create_cube_widget(self.obj, self.bones.ctrl.master, radius=0.5)
+
+
+    @classmethod
+    def add_parameters(self, params):
+        params.make_extra_control = bpy.props.BoolProperty(
+            name        = "Extra Control",
+            default     = False,
+            description = "Create an extended control"
+        )
+
+        params.make_extra_deform = bpy.props.BoolProperty(
+            name        = "Extra Deform",
+            default     = False,
+            description = "Create an extra deform bone"
+        )
+
+
+    @classmethod
+    def parameters_ui(self, layout, params):
+        r = layout.row()
+        r.prop(params, "make_extra_control", text="Extra Master Control")
+        r = layout.row()
+        r.prop(params, "make_extra_deform", text="Deform Bone")
+
+
+def create_sample(obj):
+    """ Create a sample metarig for this rig type.
+    """
+    # generated by rigify.utils.write_metarig
+    bpy.ops.object.mode_set(mode='EDIT')
+    arm = obj.data
+
+    bones = {}
+
+    bone = arm.edit_bones.new('Bone')
+    bone.head[:] = 0.0000, 0.0000, 0.0000
+    bone.tail[:] = 0.0000, 0.0000, 0.2000
+    bone.roll = 0.0000
+    bone.use_connect = False
+    bones['Bone'] = bone.name
+
+    bpy.ops.object.mode_set(mode='OBJECT')
+    pbone = obj.pose.bones[bones['Bone']]
+    pbone.rigify_type = 'basic.pivot'
+    pbone.lock_location = (False, False, False)
+    pbone.lock_rotation = (False, False, False)
+    pbone.lock_rotation_w = False
+    pbone.lock_scale = (False, False, False)
+    pbone.rotation_mode = 'QUATERNION'
+
+    bpy.ops.object.mode_set(mode='EDIT')
+    for bone in arm.edit_bones:
+        bone.select = False
+        bone.select_head = False
+        bone.select_tail = False
+    for b in bones:
+        bone = arm.edit_bones[bones[b]]
+        bone.select = True
+        bone.select_head = True
+        bone.select_tail = True
+        arm.edit_bones.active = bone
+
+    return bones
diff --git a/rigify/rigs/limbs/leg.py b/rigify/rigs/limbs/leg.py
index dbdd20cb..900af982 100644
--- a/rigify/rigs/limbs/leg.py
+++ b/rigify/rigs/limbs/leg.py
@@ -121,6 +121,14 @@ class Rig(BaseLimbRig):
 
         return name
 
+    def build_ik_pivot(self, ik_name, **args):
+        heel_bone = self.get_bone(self.bones.org.heel)
+        args = {
+            'position': (heel_bone.head + heel_bone.tail)/2,
+            **args
+        }
+        return super().build_ik_pivot(ik_name, **args)
+
     def register_switch_parents(self, pbuilder):
         super().register_switch_parents(pbuilder)
 
@@ -143,7 +151,7 @@ class Rig(BaseLimbRig):
 
     @stage.parent_bones
     def parent_heel_control_bone(self):
-        self.set_bone_parent(self.bones.ctrl.heel, self.bones.ctrl.ik)
+        self.set_bone_parent(self.bones.ctrl.heel, self.get_ik_control_output())
 
     @stage.configure_bones
     def configure_heel_control_bone(self):
@@ -200,7 +208,7 @@ class Rig(BaseLimbRig):
     @stage.parent_bones
     def parent_roll_mch_chain(self):
         chain = self.bones.mch.heel
-        self.set_bone_parent(chain[0], self.bones.ctrl.ik)
+        self.set_bone_parent(chain[0], self.get_ik_control_output())
         self.parent_bone_chain(chain)
 
     @stage.rig_bones
diff --git a/rigify/rigs/limbs/limb_rigs.py b/rigify/rigs/limbs/limb_rigs.py
index 81079c05..966feaca 100644
--- a/rigify/rigs/limbs/limb_rigs.py
+++ b/rigify/rigs/limbs/limb_rigs.py
@@ -29,6 +29,7 @@ from ...utils.naming import strip_org, make_derived_name
 from ...utils.layers import ControlLayersOption
 from ...utils.misc import pairwise_nozip, padnone, map_list
 from ...utils.switch_parent import SwitchParentBuilder
+from ...utils.components import CustomPivotControl
 
 from ...base_rig import stage, BaseRig
 
@@ -64,6 +65,7 @@ class BaseLimbRig(BaseRig):
 
         self.segments = self.params.segments
         self.bbone_segments = self.params.bbones
+        self.use_ik_pivot = self.params.make_custom_pivot
 
         rot_axis = self.params.rotation_axis
 
@@ -154,6 +156,8 @@ class BaseLimbRig(BaseRig):
     #     IK controls
     #   ik_vispole
     #     IK pole visualization.
+    #   ik_pivot
+    #     Custom IK pivot (optional).
     # mch:
     #   master:
     #     Parent of the master control.
@@ -161,6 +165,8 @@ class BaseLimbRig(BaseRig):
     #     FK follow behavior.
     #   fk[]:
     #     FK chain parents (or None)
+    #   ik_pivot
+    #     Custom IK pivot result (optional).
     #   ik_stretch
     #     IK stretch switch implementation.
     #   ik_target
@@ -332,18 +338,25 @@ class BaseLimbRig(BaseRig):
 
     def get_all_ik_controls(self):
         ctrl = self.bones.ctrl
-        return [ctrl.ik_base, ctrl.ik_pole, ctrl.ik] + self.get_extra_ik_controls()
+        controls = [ctrl.ik_base, ctrl.ik_pole, ctrl.ik]
+        if self.component_ik_pivot:
+            controls.append(self.component_ik_pivot.control)
+        return controls + self.get_extra_ik_controls()
 
     @stage.generate_bones
     def make_ik_controls(self):
         orgs = self.bones.org.main
 
-        self.bones.ctrl.ik_base = self.copy_bone(orgs[0], make_derived_name(orgs[0], 'ctrl', '_ik'))
+        self.bones.ctrl.ik_base = self.make_ik_base_bone(orgs)
         self.bones.ctrl.ik_pole = self.make_ik_pole_bone(orgs)
-        self.bones.ctrl.ik = self.make_ik_control_bone(orgs)
+        self.bones.ctrl.ik = ik_name = self.make_ik_control_bone(orgs)
 
+        self.component_ik_pivot = self.build_ik_pivot(ik_name)
         self.build_ik_parent_switch(SwitchParentBuilder(self.generator))
 
+    def make_ik_base_bone(self, orgs):
+        return self.copy_bone(orgs[0], make_derived_name(orgs[0], 'ctrl', '_ik'))
+
     def make_ik_pole_bone(self, orgs):
         name = self.copy_bone(orgs[0], make_derived_name(orgs[0], 'ctrl', '_ik_target'))
 
@@ -357,6 +370,16 @@ class BaseLimbRig(BaseRig):
     def make_ik_control_bone(self, orgs):
         return self.copy_bone(orgs[2], make_derived_name(orgs[2], 'ctrl', '_ik'))
 
+    def build_ik_pivot(self, ik_name, **args):
+        if self.use_ik_pivot:
+            return CustomPivotControl(self, 'ik_pivot', ik_name, parent=ik_name, **args)
+
+    def get_ik_control_output(self):
+        if self.component_ik_pivot:
+            return self.component_ik_pivot.output
+        else:
+            return self.bones.ctrl.ik
+
     def register_switch_parents(self, pbuilder):
         if self.rig_parent_bone:
             pbuilder.register_parent(self, self.rig_parent_bone)
@@ -453,7 +476,7 @@ class BaseLimbRig(BaseRig):
     ik_input_head_tail = 0.0
 
     def get_ik_input_bone(self):
-        return self.bones.ctrl.ik
+        return self.get_ik_control_output()
 
     def get_ik_output_chain(self):
         return [self.bones.ctrl.ik_base, self.bones.mch.ik_end, self.bones.mch.ik_target]
@@ -799,6 +822,12 @@ class BaseLimbRig(BaseRig):
             description = 'Number of segments'
         )
 
+        params.make_custom_pivot = bpy.props.BoolProperty(
+            name        = "Custom Pivot Control",
+            default     = False,
+            description = "Create a rotation pivot control that can be repositioned arbitrarily"
+        )
+
         # Setting up extra layers for the FK and tweak
         ControlLayersOption.FK.add_parameters(params)
         ControlLayersOption.TWEAK.add_parameters(params)
@@ -820,6 +849,8 @@ class BaseLimbRig(BaseRig):
         r = layout.row()
         r.prop(params, "bbones")
 
+        layout.prop(params, 'make_custom_pivot', text="Custom IK Pivot")
+
         ControlLayersOption.FK.parameters_ui(layout, params)
         ControlLayersOption.TWEAK.parameters_ui(layout, params)
 
diff --git a/rigify/rigs/limbs/paw.py b/rigify/rigs/limbs/paw.py
index f8cb1f9f..8e9d2480 100644
--- a/rigify/rigs/limbs/paw.py
+++ b/rigify/rigs/limbs/paw.py
@@ -126,7 +126,7 @@ class Rig(BaseLimbRig):
 
     @stage.parent_bones
     def parent_heel_control_bone(self):
-        self.set_bone_parent(self.bones.ctrl.heel, self.bones.ctrl.ik)
+        self.set_bone_parent(self.bones.ctrl.heel, self.get_ik_control_output())
 
     @stage.configure_bones
     def configure_heel_control_bone(self):
@@ -150,7 +150,7 @@ class Rig(BaseLimbRig):
     def parent_fk_parent_bone(self, i, parent_mch, prev_ctrl, org, prev_org):
         if i == 3:
             self.set_bone_parent(parent_mch, prev_org, use_connect=True)
-            self.set_bone_parent(self.bones.mch.toe_socket, self.bones.ctrl.ik)
+            self.set_bone_parent(self.bones.mch.toe_socket, self.get_ik_control_output())
 
         else:
             super().parent_fk_parent_bone(i, parent_mch, prev_ctrl, org, prev_org)
diff --git a/rigify/rigs/spines/basic_spine.py b/rigify/rigs/spines/basic_spine.py
index 08255bb4..00ba2f4f 100644
--- a/rigify/rigs/spines/basic_spine.py
+++ b/rigify/rigs/spines/basic_spine.py
@@ -80,13 +80,9 @@ class Rig(BaseSpineRig):
     ####################################################
     # Master control bone
 
-    @stage.generate_bones
-    def make_master_control(self):
-        super().make_master_control()
-
-        # Put the main control in the middle of the hip bone
-        base_bone = self.get_bone(self.bones.org[0])
-        put_bone(self.obj, self.bones.ctrl.master, (base_bone.head + base_bone.tail) / 2)
+    def get_master_control_pos(self, orgs):
+        base_bone = self.get_bone(orgs[0])
+        return (base_bone.head + base_bone.tail) / 2
 
     ####################################################
     # Main control bones
@@ -112,8 +108,9 @@ class Rig(BaseSpineRig):
     @stage.parent_bones
     def parent_end_control_bones(self):
         ctrl = self.bones.ctrl
-        self.set_bone_parent(ctrl.hips, ctrl.master)
-        self.set_bone_parent(ctrl.chest, ctrl.master)
+        pivot = self.get_master_control_output()
+        self.set_bone_parent(ctrl.hips, pivot)
+        self.set_bone_parent(ctrl.chest, pivot)
 
     @stage.generate_widgets
     def make_end_control_widgets(self):
@@ -247,7 +244,7 @@ class Rig(BaseSpineRig):
 
     @stage.parent_bones
     def parent_mch_chain(self):
-        master = self.bones.ctrl.master
+        master = self.get_master_control_output()
         chain = self.bones.mch.chain
         fk = self.fk_result
         for child, parent in zip(reversed(chain.hips), [master, *reversed(fk.hips)]):
diff --git a/rigify/rigs/spines/spine_rigs.py b/rigify/rigs/spines/spine_rigs.py
index 6628289f..cf563994 100644
--- a/rigify/rigs/spines/spine_rigs.py
+++ b/rigify/rigs/spines/spine_rigs.py
@@ -24,9 +24,10 @@ from itertools import count
 
 from ...utils.layers import ControlLayersOption
 from ...utils.naming import make_derived_name
-from ...utils.bones import align_bone_orientation, align_bone_to_axis
+from ...utils.bones import align_bone_orientation, align_bone_to_axis, put_bone
 from ...utils.widgets_basic import create_cube_widget
 from ...utils.switch_parent import SwitchParentBuilder
+from ...utils.components import CustomPivotControl
 
 from ...base_rig import stage
 
@@ -44,6 +45,7 @@ class BaseSpineRig(TweakChainRig):
         if len(self.bones.org) < 3:
             self.raise_error("Input to rig type must be a chain of 3 or more bones.")
 
+        self.use_torso_pivot = self.params.make_custom_pivot
         self.length = sum([self.get_bone(b).length for b in self.bones.org])
 
     ####################################################
@@ -52,6 +54,11 @@ class BaseSpineRig(TweakChainRig):
     # ctrl:
     #   master
     #     Main control.
+    #   master_pivot
+    #     Custom pivot under the master control.
+    # mch:
+    #   master_pivot
+    #     Final output of the custom pivot.
     #
     ####################################################
 
@@ -60,17 +67,42 @@ class BaseSpineRig(TweakChainRig):
 
     @stage.generate_bones
     def make_master_control(self):
-        self.bones.ctrl.master = name = self.copy_bone(self.bones.org[0], 'torso')
+        self.bones.ctrl.master = name = self.make_master_control_bone(self.bones.org)
 
+        self.component_torso_pivot = self.build_master_pivot(name)
+        self.build_parent_switch(name)
+
+    def get_master_control_pos(self, orgs):
+        return self.get_bone(orgs[0]).head
+
+    def make_master_control_bone(self, orgs):
+        name = self.copy_bone(orgs[0], 'torso')
+        put_bone(self.obj, name, self.get_master_control_pos(orgs))
         align_bone_to_axis(self.obj, name, 'y', length=self.length * 0.6)
+        return name
 
-        self.build_parent_switch(name)
+    def build_master_pivot(self, master_name, **args):
+        if self.use_torso_pivot:
+            return CustomPivotControl(
+                self, 'master_pivot', master_name, parent=master_name, **args
+            )
+
+    def get_master_control_output(self):
+        if self.component_torso_pivot:
+            return self.component_torso_pivot.output
+        else:
+            return self.bones.ctrl.master
 
     def build_parent_switch(self, master_name):
         pbuilder = SwitchParentBuilder(self.generator)
-        pbuilder.register_parent(self, master_name, name='Torso')
+
+        org_parent = self.get_bone_parent(self.bones.org[0])
+        parents = [org_parent] if org_parent else []
+
+        pbuilder.register_parent(self, self.get_master_control_output, name='Torso')
         pbuilder.build_child(
             self, master_name, exclude_self=True,
+            extra_parents=parents, select_parent=org_parent,
             prop_id='torso_parent', prop_name='Torso Parent',
             controls=lambda: self.bones.flatten('ctrl'),
         )
@@ -117,11 +149,19 @@ class BaseSpineRig(TweakChainRig):
 
     @classmethod
     def add_parameters(self, params):
+        params.make_custom_pivot = bpy.props.BoolProperty(
+            name        = "Custom Pivot Control",
+            default     = False,
+            description = "Create a rotation pivot control that can be repositioned arbitrarily"
+        )
+
         # Setting up extra layers for the FK and tweak
         ControlLayersOption.TWEAK.add_parameters(params)
 
     @classmethod
     def parameters_ui(self, layout, params):
+        layout.prop(params, 'make_custom_pivot')
+
         ControlLayersOption.TWEAK.parameters_ui(layout, params)