Cleaning up objects when armature is deleted

Rigid Body Bones/__init__.py

@@ -4,11 +4,10 @@
 # TODO support convex hull and mesh shapes
 # TODO support cone shape
 # TODO FIXED and RAGDOLL types
+# TODO different scale settings for different collision shapes
 
-# TODO cleanup when deleting the armature
 # TODO enabling/disabling bone (or changing type) and then undoing causes a hard crash
 # TODO when setting a min/max limit to 180 or -180 it disables the limit
-# TODO different scale settings for different collision shapes
 bl_info = {
     "name": "Rigid Body Bones",
     "author": "Pauan",
@@ -43,6 +42,7 @@
     properties.Bone,
 
     armatures.Update,
+    armatures.CleanupArmatures,
 
     panels.ArmaturePanel,
     panels.ArmatureSettingsPanel,

Rigid Body Bones/armatures.py

@@ -117,6 +117,14 @@ def remove_orphans(collection, exists):
     return utils.safe_remove_collection(collection)
 
 
+def remove_collection_orphans(collection, exists):
+    for sub in collection.children:
+        if sub.name not in exists:
+            utils.remove_collection_recursive(sub)
+
+    return utils.safe_remove_collection(collection)
+
+
 def make_root_body(context, armature, top):
     name = armature.data.name + " [Root]"
 
@@ -594,3 +602,33 @@ def execute(self, context):
 
 
         return {'FINISHED'}
+
+
+# This cleans up any orphan objects when an armature is deleted
+class CleanupArmatures(bpy.types.Operator):
+    bl_idname = "rigid_body_bones.cleanup_armatures"
+    bl_label = "Cleanup Rigid Body Bones"
+    # TODO use UNDO_GROUPED ?
+    bl_options = {'REGISTER', 'UNDO'}
+
+
+    @classmethod
+    def poll(cls, context):
+        return context.scene.rigid_body_bones.collection is not None
+
+
+    def execute(self, context):
+        exists = set()
+
+        for armature in bpy.data.armatures:
+            container = armature.rigid_body_bones.container
+
+            if container:
+                exists.add(container.name)
+
+        scene = context.scene.rigid_body_bones
+
+        if remove_collection_orphans(scene.collection, exists):
+            scene.property_unset("collection")
+
+        return {'FINISHED'}

Rigid Body Bones/events.py

@@ -137,6 +137,14 @@ def mode_switch():
             mark_dirty(context)
 
 
+# This is needed to cleanup the rigid body objects when the armature is deleted
+def cleanup_armatures():
+    if bpy.ops.rigid_body_bones.cleanup_armatures.poll():
+        bpy.ops.rigid_body_bones.cleanup_armatures()
+
+    return 5.0
+
+
 owner = object()
 
 def register_subscribers():
@@ -162,12 +170,17 @@ def register():
     # This is needed in order to re-subscribe when the file changes
     bpy.app.handlers.load_post.append(load_post)
 
+    bpy.app.timers.register(cleanup_armatures, persistent=True)
+
     register_subscribers()
 
 
 def unregister():
     utils.debug("UNREGISTER EVENTS")
 
+    if bpy.app.timers.is_registered(cleanup_armatures):
+        bpy.app.timers.unregister(cleanup_armatures)
+
     if bpy.app.timers.is_registered(next_tick):
         bpy.app.timers.unregister(next_tick)
 

Rigid Body Bones/utils.py

@@ -157,6 +157,18 @@ def safe_remove_collection(collection):
         return False
 
 
+def remove_collection_recursive(collection):
+    for child in collection.objects:
+        remove_object(child)
+
+    for sub in collection.children:
+        remove_collection_recursive(sub)
+
+    assert len(collection.children) == 0 and len(collection.objects) == 0
+
+    bpy.data.collections.remove(collection)
+
+
 def set_mesh_cube(mesh, dimensions):
     bm = bmesh.new()
     bmesh.ops.create_cube(bm, size=1.0, calc_uvs=False)