From 1c3cf44122f5c5f3cb82b59fca035c562c362ba9 Mon Sep 17 00:00:00 2001
From: Andrej730 <azhilenkov@gmail.com>
Date: Mon, 26 Aug 2024 17:10:07 +0500
Subject: [PATCH] get_referenced_elements to support IfcExternalInformation

---
 .../ifcopenshell/util/element.py              | 14 ++++++++++----
 .../test/util/test_element.py                 | 19 +++++++++++++++++++
 2 files changed, 29 insertions(+), 4 deletions(-)

diff --git a/src/ifcopenshell-python/ifcopenshell/util/element.py b/src/ifcopenshell-python/ifcopenshell/util/element.py
index d7af95bd17..6024cfc347 100644
--- a/src/ifcopenshell-python/ifcopenshell/util/element.py
+++ b/src/ifcopenshell-python/ifcopenshell/util/element.py
@@ -1256,13 +1256,18 @@ REFERENCE_TYPES: dict[str, ReferenceData] = {
     ),
     "IfcDocumentReference": ReferenceData("DocumentRefForObjects", "IfcRelAssociatesDocument", "RelatingDocument"),
     "IfcLibraryReference": ReferenceData("LibraryRefForObjects", "IfcRelAssociatesLibrary", "RelatingLibrary"),
+    "IfcClassification": ReferenceData(
+        "ClassificationForObjects", "IfcRelAssociatesClassification", "RelatingClassification"
+    ),
+    "IfcDocumentInformation": ReferenceData("DocumentInfoForObjects", "IfcRelAssociatesDocument", "RelatingDocument"),
+    "IfcLibraryInformation": ReferenceData("LibraryInfoForObjects", "IfcRelAssociatesLibrary", "RelatingLibrary"),
 }
 
 
 def get_referenced_elements(reference: ifcopenshell.entity_instance) -> set[ifcopenshell.entity_instance]:
     """Get all elements with assigned `reference`
 
-    :param reference: IfcExternalReference subtype reference
+    :param reference: IfcExternalReference/IfcExternalInformation subtype reference
     :type reference: ifcopenshell.entity_instance
     :return: The elements with assigned `reference`
     :rtype: set[ifcopenshell.entity_instance]
@@ -1287,9 +1292,10 @@ def get_referenced_elements(reference: ifcopenshell.entity_instance) -> set[ifco
                     related_objects.update(rel.RelatedObjects)
 
     else:
-        # IfcExternalReference
-        for external_rel in reference.ExternalReferenceForResources:
-            related_objects.update(external_rel.RelatedResourceObjects)
+        if reference.is_a("IfcExternalReference"):
+            # IfcExternalReference
+            for external_rel in reference.ExternalReferenceForResources:
+                related_objects.update(external_rel.RelatedResourceObjects)
 
         reference_data = REFERENCE_TYPES.get(ifc_class)
         if reference_data:
diff --git a/src/ifcopenshell-python/test/util/test_element.py b/src/ifcopenshell-python/test/util/test_element.py
index 864633cb15..04d8da4a64 100644
--- a/src/ifcopenshell-python/test/util/test_element.py
+++ b/src/ifcopenshell-python/test/util/test_element.py
@@ -823,6 +823,7 @@ class TestGetReferencedElements(test.bootstrap.IFC4):
             name="Foobar",
             classification=result,
         )
+        assert reference
         assert subject.get_referenced_elements(reference) == set(elements)
 
     def test_get_elements_referenced_by_library_reference(self):
@@ -843,6 +844,24 @@ class TestGetReferencedElements(test.bootstrap.IFC4):
         ifcopenshell.api.document.assign_document(self.file, document=reference, products=elements)
         assert subject.get_referenced_elements(reference) == set(elements)
 
+    def test_get_elements_referenced_by_document_info(self):
+        project = ifcopenshell.api.root.create_entity(self.file, ifc_class="IfcProject")
+        info = ifcopenshell.api.document.add_information(self.file)
+        assert subject.get_referenced_elements(info) == set([project])
+
+    def test_get_elements_referenced_by_classification(self):
+        project = ifcopenshell.api.root.create_entity(self.file, ifc_class="IfcProject")
+        info = ifcopenshell.api.classification.add_classification(self.file, classification="Name")
+        assert subject.get_referenced_elements(info) == set([project])
+
+    def test_get_elements_referenced_by_library(self):
+        info = ifcopenshell.api.library.add_library(self.file, name="Name")
+        element = ifcopenshell.api.root.create_entity(self.file, ifc_class="IfcWall")
+        rel = self.file.create_entity("IfcRelAssociatesLibrary")
+        rel.RelatingLibrary = info
+        rel.RelatedObjects = [element]
+        assert subject.get_referenced_elements(info) == set([element])
+
 
 class TestGetReferencedElementsIFC2X3(test.bootstrap.IFC2X3, TestGetReferencedElements):
     pass