Rigify: keep custom widgets already assigned in metarig.

Also make error handling more robust and extend constraint relink mixin.

5 files changed, 67 insertions, 38 deletions

diff --git a/rigify/base_generate.py b/rigify/base_generate.py
index da4949b2..16242262 100644
--- a/rigify/base_generate.py
+++ b/rigify/base_generate.py
@@ -338,7 +338,7 @@ class BaseGenerator:
         self.__run_edit_stage('prepare_bones')
 
 
-    def __auto_register_bones(self, bones, rig):
+    def __auto_register_bones(self, bones, rig, plugin=None):
         """Find bones just added and not registered by this rig."""
         for bone in bones:
             name = bone.name
@@ -347,8 +347,10 @@ class BaseGenerator:
                 if rig:
                     rig.rigify_new_bones[name] = None
 
-                if not isinstance(rig, LegacyRig):
-                    print("WARNING: rig %s didn't register bone %s\n" % (self.describe_rig(rig), name))
+                    if not isinstance(rig, LegacyRig):
+                        print("WARNING: rig %s didn't register bone %s\n" % (self.describe_rig(rig), name))
+                else:
+                    print("WARNING: plugin %s didn't register bone %s\n" % (plugin, name))
 
 
     def invoke_generate_bones(self):
@@ -365,13 +367,17 @@ class BaseGenerator:
 
             self.__auto_register_bones(self.obj.data.edit_bones, rig)
 
-        for plugin in self.plugin_list:
-            plugin.rigify_invoke_stage('generate_bones')
+        # Allow plugins to be added to the end of the list on the fly
+        for i in count(0):
+            if i >= len(self.plugin_list):
+                break
+
+            self.plugin_list[i].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, None)
+            self.__auto_register_bones(self.obj.data.edit_bones, None, plugin=self.plugin_list[i])
 
 
     def invoke_parent_bones(self):
diff --git a/rigify/rigs/basic/raw_copy.py b/rigify/rigs/basic/raw_copy.py
index 077deaa6..44c10a6c 100644
--- a/rigify/rigs/basic/raw_copy.py
+++ b/rigify/rigs/basic/raw_copy.py
@@ -47,10 +47,19 @@ class RelinkConstraintsMixin:
     def relink_bone_constraints(self, bone_name):
         if self.params.relink_constraints:
             for con in self.get_bone(bone_name).constraints:
-                parts = con.name.split('@')
+                self.relink_single_constraint(con)
 
-                if len(parts) > 1:
-                    self.relink_constraint(con, parts[1:])
+
+    relink_unmarked_constraints = False
+
+    def relink_single_constraint(self, con):
+        if self.params.relink_constraints:
+            parts = con.name.split('@')
+
+            if len(parts) > 1:
+                self.relink_constraint(con, parts[1:])
+            elif self.relink_unmarked_constraints:
+                self.relink_constraint(con, [''])
 
 
     def relink_bone_parent(self, bone_name):
@@ -73,13 +82,15 @@ class RelinkConstraintsMixin:
                 self.raise_error("Constraint {} actually has {} targets", con.name, len(con.targets))
 
             for tgt, spec in zip(con.targets, specs):
-                tgt.subtarget = self.find_relink_target(spec, tgt.subtarget)
+                if tgt.target == self.obj:
+                    tgt.subtarget = self.find_relink_target(spec, tgt.subtarget)
 
-        else:
+        elif hasattr(con, 'subtarget'):
             if len(specs) > 1:
                 self.raise_error("Only the Armature constraint can have multiple '@' targets: {}", con.name)
 
-            con.subtarget = self.find_relink_target(specs[0], con.subtarget)
+            if con.target == self.obj:
+                con.subtarget = self.find_relink_target(specs[0], con.subtarget)
 
 
     def find_relink_target(self, spec, old_target):
diff --git a/rigify/ui.py b/rigify/ui.py
index 0ed1f1a2..49c11aaf 100644
--- a/rigify/ui.py
+++ b/rigify/ui.py
@@ -608,10 +608,10 @@ class BONE_PT_rigify_buttons(bpy.types.Panel):
         if rig_name != "":
             try:
                 rig = rig_lists.rigs[rig_name]['module']
-            except (ImportError, AttributeError):
+            except (ImportError, AttributeError, KeyError):
                 row = layout.row()
                 box = row.box()
-                box.label(text="ALERT: type \"%s\" does not exist!" % rig_name)
+                box.label(text="ERROR: type \"%s\" does not exist!" % rig_name, icon='ERROR')
             else:
                 if hasattr(rig.Rig, 'parameters_ui'):
                     rig = rig.Rig
@@ -828,7 +828,7 @@ class Sample(bpy.types.Operator):
             try:
                 rig = rig_lists.rigs[self.metarig_type]["module"]
                 create_sample = rig.create_sample
-            except (ImportError, AttributeError):
+            except (ImportError, AttributeError, KeyError):
                 raise Exception("rig type '" + self.metarig_type + "' has no sample.")
             else:
                 create_sample(context.active_object)
diff --git a/rigify/utils/bones.py b/rigify/utils/bones.py
index 92b91ade..659afeae 100644
--- a/rigify/utils/bones.py
+++ b/rigify/utils/bones.py
@@ -174,7 +174,7 @@ def copy_bone(obj, bone_name, assign_name='', *, parent=False, inherit_scale=Fal
         raise MetarigError("Cannot copy bones outside of edit mode")
 
 
-def copy_bone_properties(obj, bone_name_1, bone_name_2):
+def copy_bone_properties(obj, bone_name_1, bone_name_2, transforms=True, props=True, widget=True):
     """ Copy transform and custom properties from bone 1 to bone 2. """
     if obj.mode in {'OBJECT','POSE'}:
         # Get the pose bones
@@ -182,28 +182,33 @@ def copy_bone_properties(obj, bone_name_1, bone_name_2):
         pose_bone_2 = obj.pose.bones[bone_name_2]
 
         # Copy pose bone attributes
-        pose_bone_2.rotation_mode = pose_bone_1.rotation_mode
-        pose_bone_2.rotation_axis_angle = tuple(pose_bone_1.rotation_axis_angle)
-        pose_bone_2.rotation_euler = tuple(pose_bone_1.rotation_euler)
-        pose_bone_2.rotation_quaternion = tuple(pose_bone_1.rotation_quaternion)
-
-        pose_bone_2.lock_location = tuple(pose_bone_1.lock_location)
-        pose_bone_2.lock_scale = tuple(pose_bone_1.lock_scale)
-        pose_bone_2.lock_rotation = tuple(pose_bone_1.lock_rotation)
-        pose_bone_2.lock_rotation_w = pose_bone_1.lock_rotation_w
-        pose_bone_2.lock_rotations_4d = pose_bone_1.lock_rotations_4d
+        if transforms:
+            pose_bone_2.rotation_mode = pose_bone_1.rotation_mode
+            pose_bone_2.rotation_axis_angle = tuple(pose_bone_1.rotation_axis_angle)
+            pose_bone_2.rotation_euler = tuple(pose_bone_1.rotation_euler)
+            pose_bone_2.rotation_quaternion = tuple(pose_bone_1.rotation_quaternion)
+
+            pose_bone_2.lock_location = tuple(pose_bone_1.lock_location)
+            pose_bone_2.lock_scale = tuple(pose_bone_1.lock_scale)
+            pose_bone_2.lock_rotation = tuple(pose_bone_1.lock_rotation)
+            pose_bone_2.lock_rotation_w = pose_bone_1.lock_rotation_w
+            pose_bone_2.lock_rotations_4d = pose_bone_1.lock_rotations_4d
 
         # Copy custom properties
-        for key in pose_bone_1.keys():
-            if key != "_RNA_UI" \
-            and key != "rigify_parameters" \
-            and key != "rigify_type":
-                prop1 = rna_idprop_ui_prop_get(pose_bone_1, key, create=False)
-                pose_bone_2[key] = pose_bone_1[key]
-                if prop1 is not None:
-                    prop2 = rna_idprop_ui_prop_get(pose_bone_2, key, create=True)
-                    for key in prop1.keys():
-                        prop2[key] = prop1[key]
+        if props:
+            for key in pose_bone_1.keys():
+                if key != "_RNA_UI" \
+                and key != "rigify_parameters" \
+                and key != "rigify_type":
+                    prop1 = rna_idprop_ui_prop_get(pose_bone_1, key, create=False)
+                    pose_bone_2[key] = pose_bone_1[key]
+                    if prop1 is not None:
+                        prop2 = rna_idprop_ui_prop_get(pose_bone_2, key, create=True)
+                        for key in prop1.keys():
+                            prop2[key] = prop1[key]
+
+        if widget:
+            pose_bone_2.custom_shape = pose_bone_1.custom_shape
     else:
         raise MetarigError("Cannot copy bone properties in edit mode")
 
@@ -393,9 +398,9 @@ class BoneUtilityMixin(object):
         self.register_new_bone(name, bone_name)
         return name
 
-    def copy_bone_properties(self, src_name, tgt_name):
+    def copy_bone_properties(self, src_name, tgt_name, **kwargs):
         """Copy pose-mode properties of the bone."""
-        copy_bone_properties(self.obj, src_name, tgt_name)
+        copy_bone_properties(self.obj, src_name, tgt_name, **kwargs)
 
     def rename_bone(self, old_name, new_name):
         """Rename the bone, returning the actual new name."""
diff --git a/rigify/utils/widgets.py b/rigify/utils/widgets.py
index f4713372..04e176c6 100644
--- a/rigify/utils/widgets.py
+++ b/rigify/utils/widgets.py
@@ -57,6 +57,13 @@ def obj_to_bone(obj, rig, bone_name, bone_transform_name=None):
 def create_widget(rig, bone_name, bone_transform_name=None):
     """ Creates an empty widget object for a bone, and returns the object.
     """
+    assert rig.mode != 'EDIT'
+    bone = rig.pose.bones[bone_name]
+
+    # The bone already has a widget
+    if bone.custom_shape:
+        return None
+
     obj_name = WGT_PREFIX + rig.name + '_' + bone_name
     scene = bpy.context.scene
     collection = ensure_widget_collection(bpy.context)