From 08dd8f87ff1dd6f95d307f05ea6205f57b03590d Mon Sep 17 00:00:00 2001
From: Andrej730 <azhilenkov@gmail.com>
Date: Mon, 26 Aug 2024 17:08:17 +0500
Subject: [PATCH] Operator to select document related objects

https://i.imgur.com/4gXAram.png
---
 .../bonsai/bim/module/document/__init__.py    |  1 +
 .../bonsai/bim/module/document/operator.py    | 23 +++++++++++++++++++
 src/bonsai/bonsai/bim/module/document/ui.py   |  3 +++
 3 files changed, 27 insertions(+)

diff --git a/src/bonsai/bonsai/bim/module/document/__init__.py b/src/bonsai/bonsai/bim/module/document/__init__.py
index ddbaf49985..f4eede1721 100644
--- a/src/bonsai/bonsai/bim/module/document/__init__.py
+++ b/src/bonsai/bonsai/bim/module/document/__init__.py
@@ -31,6 +31,7 @@ classes = (
     operator.LoadParentDocument,
     operator.LoadProjectDocuments,
     operator.RemoveDocument,
+    operator.SelectDocumentObjects,
     operator.UnassignDocument,
     prop.Document,
     prop.BIMDocumentProperties,
diff --git a/src/bonsai/bonsai/bim/module/document/operator.py b/src/bonsai/bonsai/bim/module/document/operator.py
index 662780c1a1..d4efe16907 100644
--- a/src/bonsai/bonsai/bim/module/document/operator.py
+++ b/src/bonsai/bonsai/bim/module/document/operator.py
@@ -20,6 +20,7 @@ import bpy
 import json
 import ifcopenshell.api
 import ifcopenshell.util.attribute
+import ifcopenshell.util.element
 import bonsai.tool as tool
 import bonsai.core.document as core
 from bonsai.bim.ifc import IfcStore
@@ -149,3 +150,25 @@ class UnassignDocument(bpy.types.Operator, tool.Ifc.Operator):
             element = tool.Ifc.get_entity(obj)
             if element:
                 core.unassign_document(tool.Ifc, product=element, document=document)
+
+
+class SelectDocumentObjects(bpy.types.Operator):
+    bl_idname = "bim.select_document_objects"
+    bl_label = "Select Document Objects"
+    bl_options = {"REGISTER", "UNDO"}
+    document: bpy.props.IntProperty(name="Document ID", default=0)
+
+    def execute(self, context):
+        if not self.document or not (relating_document := tool.Ifc.get_entity_by_id(self.document)):
+            self.report({"INFO"}, f"No document found by id '{self.document}'.")
+            return {"FINISHED"}
+
+        i = 0
+        for element in ifcopenshell.util.element.get_referenced_elements(relating_document):
+            obj = tool.Ifc.get_object(element)
+            if not obj or obj not in context.selectable_objects:
+                continue
+            obj.select_set(True)
+            i += 1
+        self.report({"INFO"}, f"{i} objects selected.")
+        return {"FINISHED"}
diff --git a/src/bonsai/bonsai/bim/module/document/ui.py b/src/bonsai/bonsai/bim/module/document/ui.py
index 573403f249..cce8043eda 100644
--- a/src/bonsai/bonsai/bim/module/document/ui.py
+++ b/src/bonsai/bonsai/bim/module/document/ui.py
@@ -66,6 +66,9 @@ class BIM_PT_documents(Panel):
             row.operator("bim.disable_editing_document", text="", icon="CANCEL")
         elif self.props.documents and self.props.active_document_index < len(self.props.documents):
             ifc_definition_id = self.props.documents[self.props.active_document_index].ifc_definition_id
+            row.operator("bim.select_document_objects", text="", icon="RESTRICT_SELECT_OFF").document = (
+                ifc_definition_id
+            )
             row.operator("bim.enable_editing_document", text="", icon="GREASEPENCIL").document = ifc_definition_id
             row.operator("bim.remove_document", text="", icon="X").document = ifc_definition_id