Merge pull request #160 from james-baber/fix-deleting-node-bug

Deleting connected bone remapping nodes now updates constraints
#159

test/unit_tests/send2ue_skeletal_meshes.py

@@ -120,6 +120,9 @@ def test_send_ue2rigify_mannequin_rig_to_unreal(self):
         bpy.data.collections['Rig'].objects.link(mannequin_rig)
         bpy.context.scene.collection.objects.unlink(mannequin_rig)
 
+        # unfreeze the rig
+        bpy.context.window_manager.ue2rigify.freeze_rig = False
+
         # set the source rig
         bpy.context.window_manager.ue2rigify.source_rig_name = 'root'
 

ue2rigify/addon/__init__.py

@@ -14,7 +14,7 @@
     "author": "Epic Games Inc.",
     "description": "Allows you to convert a given rig and its animations to a Rigify rig.",
     "blender": (2, 83, 0),
-    "version": (1, 3, 13),
+    "version": (1, 3, 14),
     "location": "3D View > Tools > UE to Rigify",
     "wiki_url": "https://epicgames.github.io/BlenderTools/ue2rigify/quickstart.html",
     "warning": "",
@@ -80,6 +80,7 @@ def register():
         bpy.utils.register_class(cls)
 
     # add an event handler that will save and load  preferences from the blend file
+    bpy.app.handlers.load_pre.append(utilities.pre_file_load)
     bpy.app.handlers.load_post.append(utilities.load_properties)
     bpy.app.handlers.save_pre.append(utilities.save_properties)
 
@@ -93,6 +94,7 @@ def unregister():
     window_manager_properties.selected_mode = window_manager_properties.source_mode
 
     # remove event handlers
+    bpy.app.handlers.load_pre.remove(utilities.pre_file_load)
     bpy.app.handlers.load_post.remove(utilities.load_properties)
     bpy.app.handlers.save_pre.remove(utilities.save_properties)
 

ue2rigify/addon/functions/nodes.py

@@ -24,12 +24,13 @@ def get_socket_names(rig_object, regex=None):
     :return list: A list of inputs.
     """
     node_socket_names = []
-    for bone in rig_object.data.bones:
-        if regex:
-            if regex.search(bone.name):
+    if rig_object:
+        for bone in rig_object.data.bones:
+            if regex:
+                if regex.search(bone.name):
+                    node_socket_names.append(bone.name)
+            else:
                 node_socket_names.append(bone.name)
-        else:
-            node_socket_names.append(bone.name)
 
     return node_socket_names
 
@@ -273,19 +274,23 @@ def remove_duplicate_socket_links(properties):
     :param object properties: The property group that contains variables that maintain the addon's correct state.
     :return bool: True or False whether or not a link was removed.
     """
+    linked_sockets = []
+    removed_sockets = False
+
     # get the node tree and sockets
     rig_node_tree = bpy.data.node_groups.get(properties.bone_tree_name)
-    linked_sockets = [(link.from_socket.name, link.to_socket.name.replace('DEF', '')) for link in rig_node_tree.links]
-    removed_sockets = False
+    if rig_node_tree:
+        if rig_node_tree.links:
+            linked_sockets = [(link.from_socket.name, link.to_socket.name.replace('DEF', '')) for link in rig_node_tree.links]
 
-    # check to see if the linked sockets already exist
-    for link in rig_node_tree.links:
-        socket_pair = link.from_socket.name, link.to_socket.name.replace('DEF', '')
+        # check to see if the linked sockets already exist
+        for link in rig_node_tree.links:
+            socket_pair = link.from_socket.name, link.to_socket.name.replace('DEF', '')
 
-        if linked_sockets.count(socket_pair) > 1:
-            rig_node_tree.links.remove(link)
-            linked_sockets.remove(socket_pair)
-            removed_sockets = True
+            if linked_sockets.count(socket_pair) > 1:
+                rig_node_tree.links.remove(link)
+                linked_sockets.remove(socket_pair)
+                removed_sockets = True
 
     return removed_sockets
 
@@ -301,29 +306,30 @@ def remove_socket_links_to_null_bones(properties):
     # get the rig objects
     control_rig_object = bpy.data.objects.get(properties.control_rig_name)
     source_rig_object = bpy.data.objects.get(properties.source_rig_name)
+    removed_sockets = False
 
     # get the node tree and sockets
     rig_node_tree = bpy.data.node_groups.get(properties.bone_tree_name)
-    linked_sockets = [(link.from_socket.name, link.to_socket.name) for link in rig_node_tree.links]
-    removed_sockets = False
+    if rig_node_tree:
+        linked_sockets = [(link.from_socket.name, link.to_socket.name) for link in rig_node_tree.links]
 
-    # check to see if the linked sockets already exist
-    for link in rig_node_tree.links:
-        socket_pair = link.from_socket.name, link.to_socket.name
+        # check to see if the linked sockets already exist
+        for link in rig_node_tree.links:
+            socket_pair = link.from_socket.name, link.to_socket.name
 
-        # check if the socket names in the links are bones that exist on the rigs
-        if control_rig_object and source_rig_object:
-            from_source_bone = source_rig_object.pose.bones.get(socket_pair[0])
-            to_source_bone = source_rig_object.pose.bones.get(socket_pair[1])
+            # check if the socket names in the links are bones that exist on the rigs
+            if control_rig_object and source_rig_object:
+                from_source_bone = source_rig_object.pose.bones.get(socket_pair[0])
+                to_source_bone = source_rig_object.pose.bones.get(socket_pair[1])
 
-            from_control_bone = control_rig_object.pose.bones.get(socket_pair[0])
-            to_control_bone = control_rig_object.pose.bones.get(socket_pair[1])
+                from_control_bone = control_rig_object.pose.bones.get(socket_pair[0])
+                to_control_bone = control_rig_object.pose.bones.get(socket_pair[1])
 
-            # if the both bones are in neither of the rigs remove the link
-            if not ((from_source_bone or to_source_bone) and (from_control_bone or to_control_bone)):
-                rig_node_tree.links.remove(link)
-                linked_sockets.remove(socket_pair)
-                removed_sockets = True
+                # if the both bones are in neither of the rigs remove the link
+                if not ((from_source_bone or to_source_bone) and (from_control_bone or to_control_bone)):
+                    rig_node_tree.links.remove(link)
+                    linked_sockets.remove(socket_pair)
+                    removed_sockets = True
 
     return removed_sockets
 
@@ -477,6 +483,29 @@ def instantiate_node(node_data, properties):
     return node
 
 
+def update_node_tree(node_tree):
+    """
+    This function updates the tracked number of nodes and links in the node tree.
+
+    :param object node_tree: A node tree object.
+    """
+    # get the properties on every update
+    properties = bpy.context.window_manager.ue2rigify
+    if remove_duplicate_socket_links(properties):
+        return None
+
+    if remove_socket_links_to_null_bones(properties):
+        return None
+
+    # if the check for updates variable is true
+    if properties.check_node_tree_for_updates:
+
+        # count the nodes and links in the tree and see if a new value should be set.
+        if properties.current_nodes_and_links != len(node_tree.links) + len(node_tree.nodes):
+            # when this property changes
+            properties.current_nodes_and_links = len(node_tree.links) + len(node_tree.nodes)
+
+
 def create_node_tree_class(classes, properties):
     """
     This function dynamically defines a node tree class from the addon's properties by subclassing type.
@@ -492,22 +521,7 @@ def update(self):
         :param class self: After this class gets dynamically defined, this becomes the traditional 'self' that is a
         reference to the class.
         """
-        # get the properties on every update
-        properties = bpy.context.window_manager.ue2rigify
-        if remove_duplicate_socket_links(properties):
-            return None
-
-        if remove_socket_links_to_null_bones(properties):
-            return None
-
-        # if the check for updates variable is true
-        if properties.check_node_tree_for_updates:
-
-            # count the nodes and links in the tree and see if a new value should be set.
-            if properties.current_nodes_and_links != len(self.links) + len(self.nodes):
-
-                # when this property changes
-                properties.current_nodes_and_links = len(self.links) + len(self.nodes)
+        update_node_tree(self)
 
     classes.append(type(
         properties.bone_tree_name.replace(' ', ''),
@@ -617,6 +631,16 @@ def init(self, context):
         for socket_name in node_class_data['inputs']:
             self.inputs.new(properties.node_socket_name.replace(' ', ''), socket_name)
 
+    def free(self):
+        """
+        This function overrides the free method in the Node class. The free method is called when the node is
+        deleted.
+
+        :param class self: After this class gets dynamically defined, this becomes the traditional 'self' that is a
+        reference to the class.
+        """
+        update_node_tree(self.rna_type.id_data)
+
     # dynamically define the node class
     node_class_definition = type(
         node_class_data['bl_idname'],
@@ -626,6 +650,7 @@ def init(self, context):
             'bl_label': node_class_data['bl_label'],
             'bl_icon': 'BONE_DATA',
             'init': init,
+            'free': free,
         }
     )
 

ue2rigify/addon/functions/scene.py

@@ -1173,7 +1173,7 @@ def edit_fk_to_source_nodes(properties):
     source_rig_object = bpy.data.objects.get(properties.source_rig_name)
     control_rig_object = create_control_rig(properties)
 
-    if control_rig_object:
+    if control_rig_object and source_rig_object:
         # constrain the fk bones to the source bones
         constrain_fk_to_source(control_rig_object, source_rig_object, properties)
 
@@ -1203,7 +1203,7 @@ def edit_source_to_deform_nodes(properties):
     source_rig_object = bpy.data.objects.get(properties.source_rig_name)
     control_rig_object = create_control_rig(properties)
 
-    if control_rig_object:
+    if control_rig_object and source_rig_object:
         # constrain the source bones to the deform bones
         constrain_source_to_deform(source_rig_object, control_rig_object, properties)
 

ue2rigify/addon/functions/utilities.py

@@ -715,10 +715,6 @@ def save_properties(*args):
             except TypeError:
                 scene_properties[attribute] = str(value)
 
-    # if the selected mode is control mode, then save the rig as frozen
-    if window_manager_properties.selected_mode == window_manager_properties.control_mode:
-        scene_properties['freeze_rig'] = True
-
 
 def load_context(properties):
     """
@@ -754,6 +750,22 @@ def load_context(properties):
                             bone[key] = value
 
 
+@bpy.app.handlers.persistent
+def pre_file_load(*args):
+    """
+    This function executes before a file load.
+
+    :param args: This soaks up the extra arguments for the app handler.
+    """
+    properties = bpy.context.window_manager.ue2rigify
+
+    # make sure that the node tree is not checking for updates
+    properties.check_node_tree_for_updates = False
+
+    # make sure that the rig is frozen
+    properties.freeze_rig = True
+
+
 @bpy.app.handlers.persistent
 def load_properties(*args):
     """
@@ -765,25 +777,32 @@ def load_properties(*args):
     window_manager_properties = bpy.context.window_manager.ue2rigify
     scene_properties = bpy.context.scene.ue2rigify
 
-    # load the freeze rig property first
-    freeze_rig_value = scene_properties.get('freeze_rig')
-    if freeze_rig_value:
-        window_manager_properties.freeze_rig = freeze_rig_value
+    # make sure that the node tree is not checking for updates
+    window_manager_properties.check_node_tree_for_updates = False
+
+    # make sure that the rig is frozen
+    window_manager_properties.freeze_rig = True
 
     # assign all the scene property values to the addon property values
     for attribute in scene_properties.keys():
         if hasattr(window_manager_properties, attribute):
-            scene_value = scene_properties.get(attribute)
-            window_manger_value = str(getattr(window_manager_properties, attribute))
+            if attribute not in ['freeze_rig', 'check_node_tree_for_updates']:
+                scene_value = scene_properties.get(attribute)
+                window_manger_value = str(getattr(window_manager_properties, attribute))
+
+                # if the scene and window manger value are not the same
+                if window_manger_value != str(scene_value):
+                    setattr(window_manager_properties, attribute, scene_value)
 
-            # if the scene and window manger value are not the same
-            if window_manger_value != str(scene_value):
-                setattr(window_manager_properties, attribute, scene_value)
+    # get the updated window manager properties
+    properties = bpy.context.window_manager.ue2rigify
+    if properties.selected_mode in [properties.fk_to_source_mode, properties.source_to_deform_mode]:
+        properties.check_node_tree_for_updates = True
 
 
 def clear_undo_history():
     """
-    This function clears blenders undo history by doing a deselect all operation and repeatedly
+    This function clears blenders undo history by calling a null operator and repeatedly
     pushing that operation into the undo stack until all previous history is cleared from the undo
     history.
     """

ue2rigify/addon/operators.py

@@ -161,6 +161,7 @@ class CreateNodesFromSelectedBones(bpy.types.Operator):
     """Create nodes that will have sockets with names of the selected bones"""
     bl_idname = "ue2rigify.create_nodes_from_selected_bones"
     bl_label = "Nodes From Selected Bones"
+    bl_options = {"REGISTER", "UNDO"}
 
     def execute(self, context):
         properties = bpy.context.window_manager.ue2rigify
@@ -172,6 +173,7 @@ class CreateLinkFromSelectedBones(bpy.types.Operator):
     """Create a pair of linked nodes from the selected bones"""
     bl_idname = "ue2rigify.create_link_from_selected_bones"
     bl_label = "Link Selected Bones"
+    bl_options = {"REGISTER", "UNDO"}
 
     def execute(self, context):
         properties = bpy.context.window_manager.ue2rigify
@@ -183,6 +185,7 @@ class CombineSelectedNodes(bpy.types.Operator):
     """Combine the selected nodes into a new node that will have the name of the active node"""
     bl_idname = "wm.combine_selected_nodes"
     bl_label = "Combine Selected Nodes"
+    bl_options = {"REGISTER", "UNDO"}
 
     @classmethod
     def poll(cls, context):
@@ -199,6 +202,7 @@ class AlignActiveNodeSockets(bpy.types.Operator):
     """Align the active node sockets with the sockets of the node it is linked to"""
     bl_idname = "wm.align_active_node_sockets"
     bl_label = "Align Active Node Sockets"
+    bl_options = {"REGISTER", "UNDO"}
 
     @classmethod
     def poll(cls, context):