From e224bfd19fd686c3f144e0cfbbe14dfea4bb9a91 Mon Sep 17 00:00:00 2001
From: Andrej730 <azhilenkov@gmail.com>
Date: Sat, 7 Sep 2024 10:47:53 +0500
Subject: [PATCH] Fix error showing material/profile psets in bbim after
 691815fd41

---
 .../ifcopenshell/util/element.py                | 17 ++++++++++++-----
 .../test/util/test_element.py                   | 11 +++++++++++
 2 files changed, 23 insertions(+), 5 deletions(-)

diff --git a/src/ifcopenshell-python/ifcopenshell/util/element.py b/src/ifcopenshell-python/ifcopenshell/util/element.py
index c666a3eea5..8804586acd 100644
--- a/src/ifcopenshell-python/ifcopenshell/util/element.py
+++ b/src/ifcopenshell-python/ifcopenshell/util/element.py
@@ -445,11 +445,18 @@ def get_elements_by_pset(pset: ifcopenshell.entity_instance) -> set[ifcopenshell
     """Retrieve the elements (or element types) that are using the provided property set."""
     is_ifc2x3 = pset.file.schema == "IFC2X3"
     elements = set()
-    rels = pset.PropertyDefinitionOf if is_ifc2x3 else pset.DefinesOccurrence
-    for rel in rels:
-        elements.update(rel.RelatedObjects)
-    for element_type in pset.DefinesType:
-        elements.add(element_type)
+    if pset.is_a("IfcPropertySet"):
+        rels = pset.PropertyDefinitionOf if is_ifc2x3 else pset.DefinesOccurrence
+        for rel in rels:
+            elements.update(rel.RelatedObjects)
+        for element_type in pset.DefinesType:
+            elements.add(element_type)
+    elif pset.is_a("IfcProfileProperties"):
+        elements.add(pset.ProfileDefinition)
+    elif pset.is_a("IfcMaterialProperties"):
+        elements.add(pset.Material)
+    else:
+        raise Exception(f"Unexpected pset type: '{pset.is_a()}' ({pset}).")
     return elements
 
 
diff --git a/src/ifcopenshell-python/test/util/test_element.py b/src/ifcopenshell-python/test/util/test_element.py
index bfeb7ddb66..759f6ff0f6 100644
--- a/src/ifcopenshell-python/test/util/test_element.py
+++ b/src/ifcopenshell-python/test/util/test_element.py
@@ -16,6 +16,7 @@
 # You should have received a copy of the GNU Lesser General Public License
 # along with IfcOpenShell.  If not, see <http://www.gnu.org/licenses/>.
 
+import ifcopenshell.api.profile
 import pytest
 import test.bootstrap
 import ifcopenshell.api.void
@@ -268,6 +269,16 @@ class TestGetElementsUsingPset(test.bootstrap.IFC4):
         ifcopenshell.api.pset.assign_pset(self.file, [element, element_type], pset)
         assert subject.get_elements_by_pset(pset) == {element, element_type}
 
+    def test_get_material_for_pset(self):
+        material = ifcopenshell.api.material.add_material(self.file)
+        pset = ifcopenshell.api.pset.add_pset(self.file, material, "FooBar")
+        assert subject.get_elements_by_pset(pset) == {material}
+
+    def test_get_profile_for_pset(self):
+        profile = ifcopenshell.api.profile.add_parameterized_profile(self.file, ifc_class="IfcRectangleProfileDef")
+        pset = ifcopenshell.api.pset.add_pset(self.file, profile, "FooBar")
+        assert subject.get_elements_by_pset(pset) == {profile}
+
 
 class TestGetElementsUsingPsetIFC2X3(test.bootstrap.IFC2X3, TestGetElementsUsingPset): ...