From 75c854da01778b82da01105d05363f3d98121cbe Mon Sep 17 00:00:00 2001
From: Andrej730 <azhilenkov@gmail.com>
Date: Mon, 15 Apr 2024 13:29:47 +0500
Subject: [PATCH] ifcopenshell.util.get_referenced_elements

---
 .../ifcopenshell/util/element.py              | 61 +++++++++++++++++++
 1 file changed, 61 insertions(+)

diff --git a/src/ifcopenshell-python/ifcopenshell/util/element.py b/src/ifcopenshell-python/ifcopenshell/util/element.py
index f5c788a8aa..5a97494737 100644
--- a/src/ifcopenshell-python/ifcopenshell/util/element.py
+++ b/src/ifcopenshell-python/ifcopenshell/util/element.py
@@ -20,6 +20,7 @@ from __future__ import annotations
 import ifcopenshell
 import ifcopenshell.util.element
 from typing import Any, Callable, Optional, Union, Literal, overload
+from collections import namedtuple
 
 
 def get_pset(
@@ -1093,6 +1094,66 @@ def get_components(element: ifcopenshell.entity_instance, include_ports=False) -
             return is_decomposed_by[0].RelatedObjects
 
 
+ReferenceData = namedtuple("ReferenceData", "inverse_attribute, rel_class, relating_element_attribute")
+
+# References below are omitted because they do not introduce
+# any additional referenced objects besides the objects
+# from their supertype IfcExternalReference
+# - IfcExternallyDefinedHatchStyle
+# - IfcExternallyDefinedSurfaceStyle
+# - IfcExternallyDefinedTextFont
+
+REFERENCE_TYPES: dict[str, ReferenceData] = {
+    "IfcClassificationReference": ReferenceData(
+        "ClassificationRefForObjects",
+        "IfcRelAssociatesClassification",
+        "RelatingClassification",
+    ),
+    "IfcDocumentReference": ReferenceData("DocumentRefForObjects", "IfcRelAssociatesDocument", "RelatingDocument"),
+    "IfcLibraryReference": ReferenceData("LibraryRefForObjects", "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
+    :type reference: ifcopenshell.entity_instance.entity_instance
+    :return: The elements with assigned `reference`
+    :rtype: set[ifcopenshell.entity_instance.entity_instance]
+
+    Example:
+
+    .. code:: python
+
+        reference = file.by_type("IfcClassificationReference")[0]
+        elements = ifcopenshell.util.element.get_referenced_elements(reference)
+    """
+
+    related_objects: set[ifcopenshell.entity_instance] = set()
+    ifc_file = reference.file
+    ifc_class = reference.is_a()
+
+    if ifc_file.schema == "IFC2X3":
+        reference_data = REFERENCE_TYPES.get(ifc_class)
+        if reference_data:
+            for rel in ifc_file.by_type(reference_data.rel_class):
+                if getattr(rel, reference_data.relating_element_attribute) == reference:
+                    related_objects.update(rel.RelatedObjects)
+
+    else:
+        # IfcExternalReference
+        for external_rel in reference.ExternalReferenceForResources:
+            related_objects.update(external_rel.RelatedResourceObjects)
+
+        reference_data = REFERENCE_TYPES.get(ifc_class)
+        if reference_data:
+            for rel in getattr(reference, reference_data.inverse_attribute):
+                related_objects.update(rel.RelatedObjects)
+
+    return related_objects
+
+
 def replace_attribute(element: ifcopenshell.entity_instance, old: Any, new: Any) -> None:
     for i, attribute_value in enumerate(element):
         if has_element_reference(attribute_value, old):