From c2de964dce9d5dbb391dc264abf46f04bff2167a Mon Sep 17 00:00:00 2001
From: Andrej730 <azhilenkov@gmail.com>
Date: Wed, 27 Mar 2024 16:59:10 +0500
Subject: [PATCH] ExtractElements - fix bug where IfcElementAssembly elements
 wasn't added to the spatial tree

Example issue - https://community.osarch.org/discussion/2054/bbim-extractelements-patch-and-missing-elements-in-ifcelementassembly
---
 .../ifcpatch/recipes/ExtractElements.py       | 38 ++++++++++++-------
 1 file changed, 24 insertions(+), 14 deletions(-)

diff --git a/src/ifcpatch/ifcpatch/recipes/ExtractElements.py b/src/ifcpatch/ifcpatch/recipes/ExtractElements.py
index 35d3e2bc08..323716bedb 100644
--- a/src/ifcpatch/ifcpatch/recipes/ExtractElements.py
+++ b/src/ifcpatch/ifcpatch/recipes/ExtractElements.py
@@ -52,8 +52,8 @@ class Patcher:
         self.query = query
 
     def patch(self):
-        self.contained_ins = {}
-        self.aggregates = {}
+        self.contained_ins: dict[str, set[ifcopenshell.entity_instance]] = {}
+        self.aggregates: dict[str, set[ifcopenshell.entity_instance]] = {}
         self.new = ifcopenshell.file(schema=self.file.wrapped_data.schema)
         self.owner_history = None
         self.reuse_identities: dict[int, ifcopenshell.entity_instance] = {}
@@ -66,18 +66,14 @@ class Patcher:
         self.create_spatial_tree()
         self.file = self.new
 
-    def add_element(self, element) -> None:
+    def add_element(self, element: ifcopenshell.entity_instance) -> None:
         new_element = self.append_asset(element)
         if not new_element:
             return
-        for rel in getattr(element, "ContainedInStructure", []):
-            spatial_element = rel.RelatingStructure
-            new_spatial_element = self.append_asset(spatial_element)
-            self.contained_ins.setdefault(spatial_element.GlobalId, set()).add(new_element)
-            self.add_decomposition_parents(spatial_element, new_spatial_element)
+        self.add_spatial_structures(element, new_element)
         self.add_decomposition_parents(element, new_element)
 
-    def append_asset(self, element) -> Union[ifcopenshell.entity_instance, None]:
+    def append_asset(self, element: ifcopenshell.entity_instance) -> Union[ifcopenshell.entity_instance, None]:
         try:
             return self.new.by_guid(element.GlobalId)
         except:
@@ -88,29 +84,43 @@ class Patcher:
             "project.append_asset", self.new, library=self.file, element=element, reuse_identities=self.reuse_identities
         )
 
-    def add_decomposition_parents(self, element, new_element) -> None:
+    def add_spatial_structures(
+        self, element: ifcopenshell.entity_instance, new_element: ifcopenshell.entity_instance
+    ) -> None:
+        """element is IfcElement"""
+        for rel in getattr(element, "ContainedInStructure", []):
+            spatial_element = rel.RelatingStructure
+            new_spatial_element = self.append_asset(spatial_element)
+            self.contained_ins.setdefault(spatial_element.GlobalId, set()).add(new_element)
+            self.add_decomposition_parents(spatial_element, new_spatial_element)
+
+    def add_decomposition_parents(
+        self, element: ifcopenshell.entity_instance, new_element: ifcopenshell.entity_instance
+    ) -> None:
+        """element is IfcObjectDefinition"""
         for rel in element.Decomposes:
             parent = rel.RelatingObject
             new_parent = self.append_asset(parent)
             self.aggregates.setdefault(parent.GlobalId, set()).add(new_element)
             self.add_decomposition_parents(parent, new_parent)
+            self.add_spatial_structures(parent, new_parent)
 
     def create_spatial_tree(self) -> None:
-        for relating_structure, related_elements in self.contained_ins.items():
+        for relating_structure_guid, related_elements in self.contained_ins.items():
             self.new.createIfcRelContainedInSpatialStructure(
                 ifcopenshell.guid.new(),
                 self.owner_history,
                 None,
                 None,
                 list(related_elements),
-                self.new.by_guid(relating_structure),
+                self.new.by_guid(relating_structure_guid),
             )
-        for relating_object, related_objects in self.aggregates.items():
+        for relating_object_guid, related_objects in self.aggregates.items():
             self.new.createIfcRelAggregates(
                 ifcopenshell.guid.new(),
                 self.owner_history,
                 None,
                 None,
-                self.new.by_guid(relating_object),
+                self.new.by_guid(relating_object_guid),
                 list(related_objects),
             )