diff --git a/src/bonsai/bonsai/bim/module/spatial/operator.py b/src/bonsai/bonsai/bim/module/spatial/operator.py index 99a2f73e6f..3ed52d9325 100644 --- a/src/bonsai/bonsai/bim/module/spatial/operator.py +++ b/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"} diff --git a/src/bonsai/bonsai/core/spatial.py b/src/bonsai/bonsai/core/spatial.py index 241e0a3df2..c70f835034 100644 --- a/src/bonsai/bonsai/core/spatial.py +++ b/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() diff --git a/src/bonsai/bonsai/tool/spatial.py b/src/bonsai/bonsai/tool/spatial.py index 20b2a4678d..f58daec3b8 100644 --- a/src/bonsai/bonsai/tool/spatial.py +++ b/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: