From 41486ba08a3ddcbd5dd40a086dd9b96ca6a3fbaa Mon Sep 17 00:00:00 2001
From: Dion Moult <dion@thinkmoult.com>
Date: Tue, 16 Nov 2021 18:07:27 +1100
Subject: [PATCH] New utility to get layers of an element. See #1869.

---
 .../ifcopenshell/util/element.py              | 16 +++++++++
 .../test/util/test_element.py                 | 36 +++++++++++++++++++
 2 files changed, 52 insertions(+)

diff --git a/src/ifcopenshell-python/ifcopenshell/util/element.py b/src/ifcopenshell-python/ifcopenshell/util/element.py
index 7faf0fd0ba..cad94b9c71 100644
--- a/src/ifcopenshell-python/ifcopenshell/util/element.py
+++ b/src/ifcopenshell-python/ifcopenshell/util/element.py
@@ -101,6 +101,22 @@ def get_material(element, should_skip_usage=False):
         return get_material(relating_type, should_skip_usage)
 
 
+def get_layers(ifc_file, element):
+    layers = []
+    representations = []
+    if getattr(element, "Representation", None):
+        representations = [element.Representation]
+    elif getattr(element, "RepresentationMaps", None):
+        representations = element.RepresentationMaps
+    for representation in representations:
+        for subelement in ifc_file.traverse(representation):
+            if subelement.is_a("IfcShapeRepresentation"):
+                layers.extend(subelement.LayerAssignments or [])
+            elif subelement.is_a("IfcGeometricRepresentationItem"):
+                layers.extend(subelement.LayerAssignment or [])
+    return layers
+
+
 def get_container(element, should_get_direct=False):
     if should_get_direct:
         if hasattr(element, "ContainedInStructure") and element.ContainedInStructure:
diff --git a/src/ifcopenshell-python/test/util/test_element.py b/src/ifcopenshell-python/test/util/test_element.py
index ca6e4cba05..a5d91996fb 100644
--- a/src/ifcopenshell-python/test/util/test_element.py
+++ b/src/ifcopenshell-python/test/util/test_element.py
@@ -214,6 +214,42 @@ class TestGetMaterial(test.bootstrap.IFC4):
         assert subject.get_material(element) == material
 
 
+class TestGetlayers(test.bootstrap.IFC4):
+    def test_getting_the_layer_of_a_product_representation(self):
+        element = ifcopenshell.api.run("root.create_entity", self.file, ifc_class="IfcWall")
+        layer = ifcopenshell.api.run("layer.add_layer", self.file)
+        representation = self.file.createIfcShapeRepresentation()
+        element.Representation = self.file.createIfcProductDefinitionShape(Representations=[representation])
+        ifcopenshell.api.run("layer.assign_layer", self.file, item=representation, layer=layer)
+        assert subject.get_layers(self.file, element) == [layer]
+
+    def test_getting_the_layer_of_a_product_item(self):
+        element = ifcopenshell.api.run("root.create_entity", self.file, ifc_class="IfcWall")
+        layer = ifcopenshell.api.run("layer.add_layer", self.file)
+        item = self.file.createIfcExtrudedAreaSolid()
+        representation = self.file.createIfcShapeRepresentation(Items=[item])
+        element.Representation = self.file.createIfcProductDefinitionShape(Representations=[representation])
+        ifcopenshell.api.run("layer.assign_layer", self.file, item=item, layer=layer)
+        assert subject.get_layers(self.file, element) == [layer]
+
+    def test_getting_the_layer_of_a_type_product_representation(self):
+        element = ifcopenshell.api.run("root.create_entity", self.file, ifc_class="IfcWallType")
+        layer = ifcopenshell.api.run("layer.add_layer", self.file)
+        representation = self.file.createIfcShapeRepresentation()
+        element.RepresentationMaps = [self.file.createIfcRepresentationMap(MappedRepresentation=representation)]
+        ifcopenshell.api.run("layer.assign_layer", self.file, item=representation, layer=layer)
+        assert subject.get_layers(self.file, element) == [layer]
+
+    def test_getting_the_layer_of_a_type_product_item(self):
+        element = ifcopenshell.api.run("root.create_entity", self.file, ifc_class="IfcWallType")
+        layer = ifcopenshell.api.run("layer.add_layer", self.file)
+        item = self.file.createIfcExtrudedAreaSolid()
+        representation = self.file.createIfcShapeRepresentation(Items=[item])
+        element.RepresentationMaps = [self.file.createIfcRepresentationMap(MappedRepresentation=representation)]
+        ifcopenshell.api.run("layer.assign_layer", self.file, item=item, layer=layer)
+        assert subject.get_layers(self.file, element) == [layer]
+
+
 class TestGetContainerIFC4(test.bootstrap.IFC4):
     def test_getting_the_spatial_container_of_an_element(self):
         element = ifcopenshell.api.run("root.create_entity", self.file, ifc_class="IfcWall")