From cf9446c4a584016d4e3fc08e8337d592a4a72200 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Bruno=20Perdig=C3=A3o?= <bruno@brunopo.com>
Date: Fri, 15 Nov 2024 08:35:30 -0300
Subject: [PATCH] Constrain part objects to aggregate object when creating
 aggregations.

---
 src/bonsai/bonsai/bim/module/aggregate/operator.py |  2 ++
 src/bonsai/bonsai/tool/aggregate.py                | 12 ++++++++++++
 2 files changed, 14 insertions(+)

diff --git a/src/bonsai/bonsai/bim/module/aggregate/operator.py b/src/bonsai/bonsai/bim/module/aggregate/operator.py
index 1bf1dbbbf6..4950d48a07 100644
--- a/src/bonsai/bonsai/bim/module/aggregate/operator.py
+++ b/src/bonsai/bonsai/bim/module/aggregate/operator.py
@@ -66,6 +66,7 @@ class BIM_OT_aggregate_assign_object(bpy.types.Operator, tool.Ifc.Operator):
                     relating_obj=relating_obj,
                     related_obj=obj,
                 )
+                tool.Aggregate.constrain_parts_to_aggregate(relating_obj)
             except core.IncompatibleAggregateError:
                 self.report({"ERROR"}, f"Cannot aggregate {obj.name} to {relating_obj.name}")
             except core.AggregateRepresentationError:
@@ -177,6 +178,7 @@ class BIM_OT_add_aggregate(bpy.types.Operator, tool.Ifc.Operator):
                     element_obj=aggregate,
                 )
             core.assign_object(tool.Ifc, tool.Aggregate, tool.Collector, relating_obj=aggregate, related_obj=obj)
+            tool.Aggregate.constrain_parts_to_aggregate(aggregate)
 
     def create_aggregate(self, context, ifc_class, aggregate_name):
         aggregate = bpy.data.objects.new(aggregate_name, None)
diff --git a/src/bonsai/bonsai/tool/aggregate.py b/src/bonsai/bonsai/tool/aggregate.py
index 9066b1ec63..bbe8e20277 100644
--- a/src/bonsai/bonsai/tool/aggregate.py
+++ b/src/bonsai/bonsai/tool/aggregate.py
@@ -83,3 +83,15 @@ class Aggregate(bonsai.core.tool.Aggregate):
             queue.update(new_parts := set(ifcopenshell.util.element.get_parts(element)))
             parts.update(new_parts)
         return parts
+
+    @classmethod
+    def constrain_parts_to_aggregate(cls, aggregate: ifcopenshell.entity_instance):
+        agg_element = tool.Ifc.get_entity(aggregate)
+        parts = ifcopenshell.util.element.get_parts(agg_element)
+        parts_objs = [tool.Ifc.get_object(p) for p in parts]
+        for obj in parts_objs:
+            constraint = next((c for c in obj.constraints if c.type == "CHILD_OF"), None)
+            if constraint:
+                obj.constraints.remove(constraint)
+            constraint = obj.constraints.new("CHILD_OF")
+            constraint.target = aggregate