Refactoring based on developer's feedback

src/bonsai/bonsai/bim/module/document/__init__.py

@@ -18,7 +18,6 @@
 
 import bpy
 from . import ui, prop, operator
-from bpy.types import VIEW3D_MT_object_context_menu
 
 classes = (
     operator.AddDocumentReference,
@@ -41,7 +40,6 @@ classes = (
     prop.Document,
     prop.DocumentObject,
     prop.AssignedDocument,
-    prop.ExpandedDocuments,
     prop.BIMDocumentProperties,
     ui.BIM_PT_documents,
     ui.BIM_PT_object_documents,
@@ -54,11 +52,9 @@ classes = (
 
 def register():
     bpy.types.Scene.BIMDocumentProperties = bpy.props.PointerProperty(type=prop.BIMDocumentProperties)
-    bpy.types.Scene.ExpandedDocuments = bpy.props.PointerProperty(type=prop.ExpandedDocuments)
-    VIEW3D_MT_object_context_menu.append(ui.add_object_documents_context_menu)
+    bpy.types.VIEW3D_MT_object_context_menu.append(ui.add_object_documents_context_menu)
 
 
 def unregister():
     del bpy.types.Scene.BIMDocumentProperties
-    del bpy.types.Scene.ExpandedDocuments
-    VIEW3D_MT_object_context_menu.remove(ui.add_object_documents_context_menu)
+    bpy.types.VIEW3D_MT_object_context_menu.remove(ui.add_object_documents_context_menu)

src/bonsai/bonsai/bim/module/document/data.py

@@ -119,13 +119,15 @@ class ObjectDocumentData:
             return results
         for rel in getattr(element, "HasAssociations", []):
             if rel.is_a("IfcRelAssociatesDocument"):
-                is_information = rel.RelatingDocument.is_a("IfcDocumentInformation")
-                is_reference = rel.RelatingDocument.is_a("IfcDocumentReference")
+                relating_document = rel.RelatingDocument
+
+                is_information = relating_document.is_a("IfcDocumentInformation")
+                is_reference = relating_document.is_a("IfcDocumentReference")
 
                 if not (is_information or is_reference):
                     continue
 
-                name = rel.RelatingDocument.Name
+                name = relating_document.Name
 
                 location = None
                 identification = None
@@ -133,33 +135,35 @@ class ObjectDocumentData:
 
                 if is_information:
                     if tool.Ifc.get_schema() == "IFC2X3":
-                        identification = rel.RelatingDocument.DocumentId
+                        identification = relating_document.DocumentId
                     else:
-                        identification = rel.RelatingDocument.Identification
+                        identification = relating_document.Identification
 
-                    location = getattr(rel.RelatingDocument, "Location", None)
+                    location = getattr(relating_document, "Location", None)
 
                 else:
-                    description = rel.RelatingDocument.Description
+                    description = relating_document.Description
                     if tool.Ifc.get_schema() == "IFC2X3":
-                        if not name and rel.RelatingDocument.ReferenceToDocument:
-                            name = rel.RelatingDocument.ReferenceToDocument[0].Name
+                        reference_to_document = relating_document.ReferenceToDocument
+                        if not name and reference_to_document:
+                            name = reference_to_document[0].Name
 
-                        identification = rel.RelatingDocument.ItemReference
-                        if not identification and rel.RelatingDocument.ReferenceToDocument:
-                            identification = rel.RelatingDocument.ReferenceToDocument[0].DocumentId
-                        location = rel.RelatingDocument.Location
+                        identification = relating_document.ItemReference
+                        if not identification and reference_to_document:
+                            identification = reference_to_document[0].DocumentId
+                        location = relating_document.Location
                     else:
-                        if not name and rel.RelatingDocument.ReferencedDocument:
-                            name = rel.RelatingDocument.ReferencedDocument.Name
+                        referenced_document = relating_document.ReferencedDocument
+                        if not name and referenced_document:
+                            name = referenced_document.Name
 
-                        identification = rel.RelatingDocument.Identification
-                        if not identification and rel.RelatingDocument.ReferencedDocument:
-                            identification = rel.RelatingDocument.ReferencedDocument.Identification
+                        identification = relating_document.Identification
+                        if not identification and referenced_document:
+                            identification = referenced_document.Identification
 
-                        location = rel.RelatingDocument.Location
-                        if location is None and rel.RelatingDocument.ReferencedDocument:
-                            location = rel.RelatingDocument.ReferencedDocument.Location
+                        location = relating_document.Location
+                        if location is None and referenced_document:
+                            location = referenced_document.Location
 
                 if location:
                     if not "://" in location:
@@ -169,7 +173,7 @@ class ObjectDocumentData:
 
                 results.append(
                     {
-                        "id": rel.RelatingDocument.id(),
+                        "id": relating_document.id(),
                         "identification": identification,
                         "name": name,
                         "location": location,

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

@@ -35,10 +35,8 @@ def update_document_objects(document_id=None):
 
     if document_id is None:
         props = tool.Document.get_document_props()
-        if props.documents and props.active_document_index < len(props.documents):
-            document = props.documents[props.active_document_index]
-            if document.ifc_definition_id:
-                document_id = document.ifc_definition_id
+        if props.active_document and props.active_document.ifc_definition_id:
+            document_id = props.active_document.ifc_definition_id
 
     if document_id:
         DocumentData.load_document_objects_into_props(document_id)
@@ -122,8 +120,8 @@ class AddInformation(bpy.types.Operator, tool.Ifc.Operator):
     def _execute(self, context):
         props = tool.Document.get_document_props()
         parent = None
-        if props.documents and props.active_document_index < len(props.documents):
-            selected_document = props.documents[props.active_document_index]
+        if props.active_document:
+            selected_document = props.active_document
 
             if selected_document.ifc_definition_id == -1:
                 parent = tool.Ifc.get().by_type("IfcProject")[0] if tool.Ifc.get().by_type("IfcProject") else None
@@ -139,7 +137,7 @@ class AddInformation(bpy.types.Operator, tool.Ifc.Operator):
 
         expanded_docs = []
         try:
-            expanded_docs = json.loads(context.scene.ExpandedDocuments.json_string)
+            expanded_docs = json.loads(props.json_string)
         except (AttributeError, json.JSONDecodeError):
             pass
 
@@ -152,7 +150,7 @@ class AddInformation(bpy.types.Operator, tool.Ifc.Operator):
             if parent.id() not in expanded_docs:
                 expanded_docs.append(parent.id())
 
-        context.scene.ExpandedDocuments.json_string = json.dumps(expanded_docs)
+        props.json_string = json.dumps(expanded_docs)
 
         bpy.ops.bim.load_project_documents()
 
@@ -165,11 +163,11 @@ class AddDocumentReference(bpy.types.Operator, tool.Ifc.Operator):
     def _execute(self, context):
         props = tool.Document.get_document_props()
 
-        if not props.documents or props.active_document_index >= len(props.documents):
+        if not props.active_document:
             self.report({"ERROR"}, "No document selected")
             return {"CANCELLED"}
 
-        selected_document = props.documents[props.active_document_index]
+        selected_document = props.active_document
 
         if not selected_document.is_information:
             self.report({"ERROR"}, "Cannot add a reference to a reference element")
@@ -181,13 +179,13 @@ class AddDocumentReference(bpy.types.Operator, tool.Ifc.Operator):
         core.add_reference(tool.Ifc, tool.Document)
         expanded_docs = []
         try:
-            expanded_docs = json.loads(context.scene.ExpandedDocuments.json_string)
+            expanded_docs = json.loads(props.json_string)
         except (AttributeError, json.JSONDecodeError):
             pass
 
         if parent.id() not in expanded_docs:
             expanded_docs.append(parent.id())
-            context.scene.ExpandedDocuments.json_string = json.dumps(expanded_docs)
+            props.json_string = json.dumps(expanded_docs)
 
         bpy.ops.bim.load_project_documents()
 
@@ -292,9 +290,8 @@ class UnassignDocument(bpy.types.Operator, tool.Ifc.Operator):
                 core.unassign_document(tool.Ifc, product=element, document=document)
         props = tool.Document.get_document_props()
         active_document_id = None
-        if props.documents and props.active_document_index < len(props.documents):
-            active_document = props.documents[props.active_document_index]
-            active_document_id = active_document.ifc_definition_id
+        if props.active_document:
+            active_document_id = props.active_document.ifc_definition_id
 
         if active_document_id and active_document_id != self.document:
             update_document_objects(active_document_id)
@@ -420,8 +417,9 @@ class ToggleDocument(bpy.types.Operator, tool.Ifc.Operator):
 
     def _execute(self, context):
         expanded_documents = []
+        props = tool.Document.get_document_props()
         try:
-            expanded_documents = json.loads(context.scene.ExpandedDocuments.json_string)
+            expanded_documents = json.loads(props.json_string)
         except (AttributeError, json.JSONDecodeError):
             expanded_documents = []
 
@@ -439,8 +437,7 @@ class ToggleDocument(bpy.types.Operator, tool.Ifc.Operator):
                 expanded_documents.append(virtual_root_id)
             elif self.option == "Collapse" and virtual_root_id in expanded_documents:
                 expanded_documents.remove(virtual_root_id)
-
-        context.scene.ExpandedDocuments.json_string = json.dumps(expanded_documents)
+        props.json_string = json.dumps(expanded_documents)
 
         bpy.ops.bim.load_project_documents()
         return {"FINISHED"}

src/bonsai/bonsai/bim/module/document/prop.py

@@ -80,13 +80,6 @@ class Document(PropertyGroup):
         is_expanded: bool
 
 
-class ExpandedDocuments(PropertyGroup):
-    json_string: StringProperty(name="JSON String", default="[]")
-
-    if TYPE_CHECKING:
-        json_string: str
-
-
 class DocumentObject(PropertyGroup):
     name: StringProperty(name="Name")
     ifc_definition_id: IntProperty(name="IFC Definition ID")
@@ -124,6 +117,7 @@ class BIMDocumentProperties(PropertyGroup):
     active_document_object_index: IntProperty(name="Active Document Object Index")
     assigned_documents: CollectionProperty(name="Assigned Documents", type=AssignedDocument)
     active_assigned_document_index: IntProperty(name="Active Assigned Document Index")
+    json_string: StringProperty(name="JSON String", default="[]")
 
     if TYPE_CHECKING:
         document_attributes: bpy.types.bpy_prop_collection_idprop[Attribute]
@@ -137,6 +131,7 @@ class BIMDocumentProperties(PropertyGroup):
         active_document_object_index: int
         assigned_documents: bpy.types.bpy_prop_collection_idprop[AssignedDocument]
         active_assigned_document_index: int
+        json_string: str
 
     @property
     def active_document(self) -> Union[Document, None]:

src/bonsai/bonsai/bim/module/document/ui.py

@@ -68,14 +68,11 @@ class BIM_PT_documents(Panel):
             row.operator("bim.edit_document", text="", icon="CHECKMARK")
             row.operator("bim.disable_editing_document", text="", icon="CANCEL")
         else:
-            if not self.props.documents or not self.props.active_document_index < len(self.props.documents) or \
-            (self.props.active_document_index < len(self.props.documents) and 
-                self.props.documents[self.props.active_document_index].is_information):
+            if not self.props.active_document or self.props.active_document.is_information:
                 row.operator("bim.add_information", text="", icon="ADD")
 
-            if self.props.documents and self.props.active_document_index < len(self.props.documents):
-                active_doc = self.props.documents[self.props.active_document_index]
-                if active_doc.is_information and active_doc.ifc_definition_id != -1:
+            if self.props.active_document:
+                if self.props.active_document.is_information and self.props.active_document.ifc_definition_id != -1:
                     row.operator("bim.add_document_reference", text="", icon="FILE_HIDDEN")
 
             active_document = self.props.active_document
@@ -84,7 +81,7 @@ class BIM_PT_documents(Panel):
                 row.operator("bim.select_document_objects", text="", icon="RESTRICT_SELECT_OFF").document = (
                     ifc_definition_id
                 )
-                
+
                 row.operator("bim.assign_document", text="", icon="BRUSH_DATA").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
@@ -97,12 +94,8 @@ class BIM_PT_documents(Panel):
             else:
                 draw_attributes(self.props.document_attributes, self.layout, filter_attributes=["Name"])
 
-        if (
-            self.props.is_editing
-            and self.props.documents
-            and self.props.active_document_index < len(self.props.documents)
-        ):
-            document = self.props.documents[self.props.active_document_index]
+        if self.props.is_editing and self.props.active_document:
+            document = self.props.active_document
             box = self.layout.box()
             row = box.row(align=True)
             row.label(text="Assigned Objects", icon="OUTLINER_OB_EMPTY")
@@ -181,21 +174,19 @@ class BIM_PT_object_documents(Panel):
             row = self.layout.row(align=True)
             row.alignment = "RIGHT"
 
-            if self.props.documents and self.props.active_document_index < len(self.props.documents):
-                document = self.props.documents[self.props.active_document_index]
+            if self.props.active_document:
+                document = self.props.active_document
 
                 assigned_doc_ids = []
                 for doc in ObjectDocumentData.data["documents"]:
                     assigned_doc_ids.append(doc["id"])
 
                 # Only show assign button if the document is information (not reference) and not already assigned
-                if (document.is_information and 
-                    document.ifc_definition_id not in assigned_doc_ids):
+                if document.is_information and document.ifc_definition_id not in assigned_doc_ids:
                     doc_op = row.operator("bim.assign_document", text="", icon="BRUSH_DATA")
                     doc_op.document = document.ifc_definition_id  # Pass the current document's ID
                 elif document.ifc_definition_id in assigned_doc_ids:
                     row.label(text="", icon="CHECKMARK")
-
             self.layout.template_list(
                 "BIM_UL_documents", "", self.props, "documents", self.props, "active_document_index"
             )
@@ -249,8 +240,8 @@ class BIM_UL_document_objects(UIList):
             row.operator("bim.select_object", text="", icon="RESTRICT_SELECT_OFF").obj_name = item.name
 
             props = tool.Document.get_document_props()
-            if props.documents and props.active_document_index < len(props.documents):
-                document = props.documents[props.active_document_index]
+            if props.active_document:
+                document = props.active_document
 
                 op = row.operator("bim.unassign_document", text="", icon="X")
                 op.document = document.ifc_definition_id

src/bonsai/bonsai/core/document.py

@@ -18,8 +18,6 @@
 
 from __future__ import annotations
 from typing import TYPE_CHECKING, Optional
-import bpy
-import json
 
 if TYPE_CHECKING:
     import ifcopenshell
@@ -34,15 +32,7 @@ def load_project_documents(document: tool.Document) -> None:
 
 def load_document(document_tool: tool.Document, document: ifcopenshell.entity_instance) -> None:
     document_tool.clear_document_tree()
-    try:
-        expanded_docs = json.loads(bpy.context.scene.ExpandedDocuments.json_string)
-    except (AttributeError, json.JSONDecodeError):
-        expanded_docs = []
-
-    if document.id() not in expanded_docs:
-        expanded_docs.append(document.id())
-        bpy.context.scene.ExpandedDocuments.json_string = json.dumps(expanded_docs)
-
+    document_tool.expand_document(document)
     document_tool.import_project_documents()
     document_tool.disable_editing_document()
 
@@ -53,60 +43,40 @@ def disable_document_editing_ui(document: tool.Document) -> None:
 
 
 def enable_editing_document(document_tool: tool.Document, document: ifcopenshell.entity_instance) -> None:
-    props = document_tool.get_document_props()
-    props.active_document_id = document.id()
-    props.is_document_editing = True
+    document_tool.set_active_document(document)
+    document_tool.enable_document_editing()
     document_tool.import_document_attributes(document)
 
 
 def disable_editing_document(document: tool.Document) -> None:
-    props = document.get_document_props()
-    props.active_document_id = 0
-    props.is_document_editing = False
-    props.document_attributes.clear()
+    document.clear_active_document()
+    document.disable_document_editing()
+    document.clear_document_attributes()
 
 
 def add_information(ifc: tool.Ifc, document_tool: tool.Document, parent=None) -> ifcopenshell.entity_instance:
     document_tool.clear_document_tree()
 
-    if parent is None and ifc.get().by_type("IfcProject"):
-        parent = ifc.get().by_type("IfcProject")[0]
+    if parent is None:
+        parent = document_tool.get_default_parent_for_information(ifc)
 
     information = ifc.run("document.add_information", parent=parent)
     ifc.run("document.add_reference", information=information)
-    if parent and parent.is_a("IfcDocumentInformation"):
-        try:
-            expanded_docs = json.loads(bpy.context.scene.ExpandedDocuments.json_string)
-        except (AttributeError, json.JSONDecodeError):
-            expanded_docs = []
 
-        if parent.id() not in expanded_docs:
-            expanded_docs.append(parent.id())
-            bpy.context.scene.ExpandedDocuments.json_string = json.dumps(expanded_docs)
+    if document_tool.is_document_information(parent):
+        document_tool.expand_document(parent)
 
     document_tool.import_project_documents()
+    return information
 
 
 def add_reference(ifc: tool.Ifc, document: tool.Document) -> None:
-    props = document.get_document_props()
-    parent = None
-
-    if props.documents and props.active_document_index < len(props.documents):
-        selected_document = props.documents[props.active_document_index]
-        if selected_document.is_information:
-            parent = ifc.get().by_id(selected_document.ifc_definition_id)
+    parent = document.get_selected_document_information(ifc)
 
     if parent:
         reference = ifc.run("document.add_reference", information=parent)
         reference.Location = ""
-        try:
-            expanded_docs = json.loads(bpy.context.scene.ExpandedDocuments.json_string)
-        except (AttributeError, json.JSONDecodeError):
-            expanded_docs = []
-
-        if parent.id() not in expanded_docs:
-            expanded_docs.append(parent.id())
-            bpy.context.scene.ExpandedDocuments.json_string = json.dumps(expanded_docs)
+        document.expand_document(parent)
 
     document.import_project_documents()
 

src/bonsai/bonsai/tool/document.py

@@ -23,6 +23,7 @@ import ifcopenshell.util.system
 import bonsai.bim.helper
 import bonsai.core.tool
 import bonsai.tool as tool
+import json
 from typing import Any, Union, TYPE_CHECKING
 
 if TYPE_CHECKING:
@@ -86,7 +87,7 @@ class Document(bonsai.core.tool.Document):
         props.documents.clear()
         file = tool.Ifc.get()
         try:
-            expanded_documents = json.loads(bpy.context.scene.ExpandedDocuments.json_string)
+            expanded_documents = json.loads(props.json_string)
         except (AttributeError, json.JSONDecodeError):
             expanded_documents = []
 
@@ -116,7 +117,7 @@ class Document(bonsai.core.tool.Document):
                     document_children[parent_id].append(ref)
         else:
             for ref in file.by_type("IfcDocumentReference"):
-                if hasattr(ref, "ReferencedDocument") and ref.ReferencedDocument:
+                if ref.ReferencedDocument:
                     parent = ref.ReferencedDocument
                     parent_id = parent.id()
                     if parent_id not in document_children:
@@ -198,21 +199,18 @@ class Document(bonsai.core.tool.Document):
 
             info_children = [d for d in children if d.is_a("IfcDocumentInformation")]
             ref_children = [d for d in children if not d.is_a("IfcDocumentInformation")]
-            
+
             info_children.sort(
-                key=lambda doc: (
-                    (cls.get_document_information_id(doc) or "").lower(),
-                    (doc.Name or "").lower()
-                )
+                key=lambda doc: ((cls.get_document_information_id(doc) or "").lower(), (doc.Name or "").lower())
             )
-            
+
             ref_children.sort(
                 key=lambda doc: (
                     (cls.get_external_reference_id(doc) or "").lower(),
-                    (doc.Description or doc.Name or "").lower()
+                    (doc.Description or doc.Name or "").lower(),
                 )
             )
-            
+
             for child in info_children + ref_children:
                 cls._process_document(child, props, document_children, expanded_documents, depth + 1)
 
@@ -253,3 +251,48 @@ class Document(bonsai.core.tool.Document):
         if document.file.schema == "IFC2X3":
             return document.DocumentReferences or ()
         return document.HasDocumentReferences
+
+    @classmethod
+    def enable_document_editing(cls) -> None:
+        props = cls.get_document_props()
+        props.is_editing = True
+
+    @classmethod
+    def disable_document_editing(cls) -> None:
+        props = cls.get_document_props()
+        props.is_editing = False
+
+    @classmethod
+    def clear_active_document(cls) -> None:
+        props = cls.get_document_props()
+        props.active_document_id = 0
+
+    @classmethod
+    def clear_document_attributes(cls) -> None:
+        props = cls.get_document_props()
+        props.document_attributes.clear()
+
+    @classmethod
+    def expand_document(cls, document: ifcopenshell.entity_instance) -> None:
+        props = cls.get_document_props()
+        try:
+            expanded_docs = json.loads(props.json_string)
+        except (AttributeError, json.JSONDecodeError):
+            expanded_docs = []
+
+        if document.id() not in expanded_docs:
+            expanded_docs.append(document.id())
+            props.json_string = json.dumps(expanded_docs)
+
+    @classmethod
+    def get_default_parent_for_information(cls, ifc) -> Union[ifcopenshell.entity_instance, None]:
+        projects = ifc.get().by_type("IfcProject")
+        return projects[0] if projects else None
+
+    @classmethod
+    def get_selected_document_information(cls, ifc) -> Union[ifcopenshell.entity_instance, None]:
+        props = cls.get_document_props()
+
+        if props.active_document and props.active_document.is_information:
+            return ifc.get().by_id(props.active_document.ifc_definition_id)
+        return None