Implement recursive contract / expand of the spatial tree

src/bonsai/bonsai/bim/module/spatial/operator.py

@@ -336,22 +336,40 @@ class ImportSpatialDecomposition(bpy.types.Operator):
 class ContractContainer(bpy.types.Operator):
     bl_idname = "bim.contract_container"
     bl_label = "Contract Container"
+    bl_description = "Contract the hierarchy\nALT+CLICK to recursively contract"
     bl_options = {"REGISTER", "UNDO"}
     container: bpy.props.IntProperty()
+    is_recursive: bpy.props.BoolProperty(name="Is Recursive", default=False, options={"SKIP_SAVE"})
+
+    def invoke(self, context, event):
+        if event.type == "LEFTMOUSE" and event.alt:
+            self.is_recursive = True
+        return self.execute(context)
 
     def execute(self, context):
-        core.contract_container(tool.Spatial, container=tool.Ifc.get().by_id(self.container))
+        core.contract_container(
+            tool.Spatial, container=tool.Ifc.get().by_id(self.container), is_recursive=self.is_recursive
+        )
         return {"FINISHED"}
 
 
 class ExpandContainer(bpy.types.Operator):
     bl_idname = "bim.expand_container"
     bl_label = "Expand Container"
+    bl_description = "Expand the hierarchy\nALT+CLICK to recursively contract"
     bl_options = {"REGISTER", "UNDO"}
     container: bpy.props.IntProperty()
+    is_recursive: bpy.props.BoolProperty(name="Is Recursive", default=False, options={"SKIP_SAVE"})
+
+    def invoke(self, context, event):
+        if event.type == "LEFTMOUSE" and event.alt:
+            self.is_recursive = True
+        return self.execute(context)
 
     def execute(self, context):
-        core.expand_container(tool.Spatial, container=tool.Ifc.get().by_id(self.container))
+        core.expand_container(
+            tool.Spatial, container=tool.Ifc.get().by_id(self.container), is_recursive=self.is_recursive
+        )
         return {"FINISHED"}
 
 

src/bonsai/bonsai/core/spatial.py

@@ -125,13 +125,13 @@ def set_orientation_slot(spatial: tool.Spatial, container: ifcopenshell.entity_i
     spatial.create_orientation_slot(container)
 
 
-def contract_container(spatial: tool.Spatial, container: ifcopenshell.entity_instance) -> None:
-    spatial.contract_container(container)
+def contract_container(spatial: tool.Spatial, container: ifcopenshell.entity_instance, is_recursive: bool) -> None:
+    spatial.contract_container(container, is_recursive=is_recursive)
     spatial.import_spatial_decomposition()
 
 
-def expand_container(spatial: tool.Spatial, container: ifcopenshell.entity_instance) -> None:
-    spatial.expand_container(container)
+def expand_container(spatial: tool.Spatial, container: ifcopenshell.entity_instance, is_recursive: bool) -> None:
+    spatial.expand_container(container, is_recursive=is_recursive)
     spatial.import_spatial_decomposition()
 
 

src/bonsai/bonsai/tool/spatial.py

@@ -558,18 +558,28 @@ class Spatial(bonsai.core.tool.Spatial):
             return container
 
     @classmethod
-    def contract_container(cls, container: ifcopenshell.entity_instance) -> None:
+    def contract_container(cls, container: ifcopenshell.entity_instance, is_recursive: bool) -> None:
         props = cls.get_spatial_props()
-        contracted_containers = json.loads(props.contracted_containers)
-        contracted_containers.append(container.id())
-        props.contracted_containers = json.dumps(contracted_containers)
+        contracted_containers = set(json.loads(props.contracted_containers))
+        queue = [container]
+        while queue:
+            item = queue.pop()
+            if is_recursive and (children := ifcopenshell.util.element.get_parts(item)):
+                queue.extend(children)
+            contracted_containers.add(item.id())
+        props.contracted_containers = json.dumps(list(contracted_containers))
 
     @classmethod
-    def expand_container(cls, container: ifcopenshell.entity_instance) -> None:
+    def expand_container(cls, container: ifcopenshell.entity_instance, is_recursive: bool) -> None:
         props = cls.get_spatial_props()
-        contracted_containers = json.loads(props.contracted_containers)
-        contracted_containers.remove(container.id())
-        props.contracted_containers = json.dumps(contracted_containers)
+        contracted_containers = set(json.loads(props.contracted_containers))
+        queue = [container]
+        while queue:
+            item = queue.pop()
+            if is_recursive and (children := ifcopenshell.util.element.get_parts(item)):
+                queue.extend(children)
+            contracted_containers.discard(item.id())
+        props.contracted_containers = json.dumps(list(contracted_containers))
 
     @classmethod
     def toggle_container_element(cls, element_index: int, is_recursive: bool) -> None: