From f364643d7e2d7c06534cb953656574b0dd621f0d Mon Sep 17 00:00:00 2001
From: Andrej730 <azhilenkov@gmail.com>
Date: Tue, 25 Jun 2024 17:46:42 +0500
Subject: [PATCH] fix assigning material usages / reloading reprs after
 c4ae9578e #4843

---
 .../blenderbim/bim/module/model/product.py       | 16 ++++++++++++++++
 src/blenderbim/blenderbim/tool/geometry.py       |  4 ++++
 src/blenderbim/blenderbim/tool/model.py          |  3 +--
 .../ifcopenshell/api/material/assign_material.py |  2 +-
 .../ifcopenshell/util/element.py                 | 11 +++++++++++
 5 files changed, 33 insertions(+), 3 deletions(-)

diff --git a/src/blenderbim/blenderbim/bim/module/model/product.py b/src/blenderbim/blenderbim/bim/module/model/product.py
index 0da8d17170..36dcba4f26 100644
--- a/src/blenderbim/blenderbim/bim/module/model/product.py
+++ b/src/blenderbim/blenderbim/bim/module/model/product.py
@@ -515,6 +515,22 @@ def regenerate_profile_usage(usecase_path, ifc_file, settings):
 
 def ensure_material_assigned(usecase_path: str, ifc_file: ifcopenshell.file, settings: dict[str, Any]) -> None:
     elements = settings["products"]
+    material = settings.get("material")
+    if material:
+        assigned_material = settings["material"]
+    else:
+        material_type: ifcopenshell.util.element.MATERIAL_TYPE = settings["type"]
+        element = elements[0]
+        if material_type == "IfcMaterial":
+            assigned_material = ifcopenshell.util.element.get_material(element, should_inherit=False)
+            assert assigned_material  # Type checker.
+        # 1) Material usages just inherit the style from the type material, so can't override it.
+        # 2) If type is Set and no material argument were provided, then Set was just created
+        # and not yet have any IfcMaterials.
+        elif material_type.endswith("Usage") or material_type.endswith("Set"):
+            return
+        elif material_type == "IfcMaterialList":
+            assert False, "Current assign_material implementation requires 'material' argument for IfcMaterialList."
 
     for element in elements[:]:
         if element.is_a("IfcElementType"):
diff --git a/src/blenderbim/blenderbim/tool/geometry.py b/src/blenderbim/blenderbim/tool/geometry.py
index 95fcd142e4..b12cc9a750 100644
--- a/src/blenderbim/blenderbim/tool/geometry.py
+++ b/src/blenderbim/blenderbim/tool/geometry.py
@@ -526,6 +526,10 @@ class Geometry(blenderbim.core.tool.Geometry):
             return False
         own_material = ifcopenshell.util.element.get_material(element, should_inherit=False)
         if own_material:
+            # Material usages just inherit the style from the type material, so can't override it.
+            if own_material.is_a("IfcMaterialUsageDefinition"):
+                return False
+            own_material = ifcopenshell.util.element.get_materials(element, should_inherit=False)[0]
             inherited_style = cls.get_inherited_material_style(element)
             style = tool.Material.get_style(own_material) if own_material else None
             if inherited_style != style:
diff --git a/src/blenderbim/blenderbim/tool/model.py b/src/blenderbim/blenderbim/tool/model.py
index d7e221fd7c..4225b09acf 100644
--- a/src/blenderbim/blenderbim/tool/model.py
+++ b/src/blenderbim/blenderbim/tool/model.py
@@ -1268,7 +1268,6 @@ class Model(blenderbim.core.tool.Model):
         after material assignment or material unassignment.
         """
 
-        style_object = None
         if assigned_material:
             # NOTE: currently only IfcMaterials are supported
             # for anyone else we just switch representation.
@@ -1340,7 +1339,7 @@ class Model(blenderbim.core.tool.Model):
         occurrences = [
             e
             for e in ifcopenshell.util.element.get_types(element_type)
-            if not ifcopenshell.util.element.get_material(e, should_inherit=False)
+            if not tool.Geometry.has_material_style_override(e)
         ]
         return occurrences
 
diff --git a/src/ifcopenshell-python/ifcopenshell/api/material/assign_material.py b/src/ifcopenshell-python/ifcopenshell/api/material/assign_material.py
index a238d7a8c9..5ddf9c9524 100644
--- a/src/ifcopenshell-python/ifcopenshell/api/material/assign_material.py
+++ b/src/ifcopenshell-python/ifcopenshell/api/material/assign_material.py
@@ -27,7 +27,7 @@ from typing import Optional, Union
 def assign_material(
     file: ifcopenshell.file,
     products: list[ifcopenshell.entity_instance],
-    type: str = "IfcMaterial",
+    type: ifcopenshell.util.element.MATERIAL_TYPE = "IfcMaterial",
     material: Optional[ifcopenshell.entity_instance] = None,
 ) -> Union[ifcopenshell.entity_instance, list[ifcopenshell.entity_instance], None]:
     """Assigns a material to the list of products
diff --git a/src/ifcopenshell-python/ifcopenshell/util/element.py b/src/ifcopenshell-python/ifcopenshell/util/element.py
index fb76b72f89..e9aeb99403 100644
--- a/src/ifcopenshell-python/ifcopenshell/util/element.py
+++ b/src/ifcopenshell-python/ifcopenshell/util/element.py
@@ -23,6 +23,17 @@ from typing import Any, Callable, Optional, Union, Literal, overload
 from collections import namedtuple
 
 
+MATERIAL_TYPE = Literal[
+    "IfcMaterial",
+    "IfcMaterialConstituentSet",
+    "IfcMaterialLayerSet",
+    "IfcMaterialLayerSetUsage",
+    "IfcMaterialProfileSet",
+    "IfcMaterialProfileSetUsage",
+    "IfcMaterialList",
+]
+
+
 def get_pset(
     element: ifcopenshell.entity_instance,
     name: str,