From fb6c670aa6770d95abdfc4db4b4515c6930d703c Mon Sep 17 00:00:00 2001
From: Dion Moult <dion@thinkmoult.com>
Date: Wed, 7 Jul 2021 10:25:19 +1000
Subject: [PATCH] Fix colour by undo operators in search module. See #1475.

---
 .../blenderbim/bim/module/search/operator.py  | 78 +++++++++++++++++--
 1 file changed, 72 insertions(+), 6 deletions(-)

diff --git a/src/blenderbim/blenderbim/bim/module/search/operator.py b/src/blenderbim/blenderbim/bim/module/search/operator.py
index d5a9438c08..f4b938d654 100644
--- a/src/blenderbim/blenderbim/bim/module/search/operator.py
+++ b/src/blenderbim/blenderbim/bim/module/search/operator.py
@@ -137,6 +137,14 @@ class ColourByAttribute(bpy.types.Operator):
     bl_options = {"REGISTER", "UNDO"}
 
     def execute(self, context):
+        IfcStore.begin_transaction(self)
+        self.store_state()
+        result = self._execute(context)
+        IfcStore.add_transaction_operation(self)
+        IfcStore.end_transaction(self)
+        return result
+
+    def _execute(self, context):
         self.file = IfcStore.get_file()
         colours = cycle(colour_list)
         values = {}
@@ -153,10 +161,24 @@ class ColourByAttribute(bpy.types.Operator):
             if value not in values:
                 values[value] = next(colours)
             obj.color = values[value]
-        area = next(area for area in context.screen.areas if area.type == "VIEW_3D")
-        area.spaces[0].shading.color_type = "OBJECT"
+        areas = [a for a in bpy.context.screen.areas if a.type == "VIEW_3D"]
+        if areas:
+            areas[0].spaces[0].shading.color_type = "OBJECT"
         return {"FINISHED"}
 
+    def store_state(self):
+        areas = [a for a in bpy.context.screen.areas if a.type == "VIEW_3D"]
+        if areas:
+            self.transaction_data = {"area": areas[0], "color_type": areas[0].spaces[0].shading.color_type}
+
+    def rollback(self, data):
+        if data:
+            data["area"].spaces[0].shading.color_type = data["color_type"]
+
+    def commit(self, data):
+        if data:
+            data["area"].spaces[0].shading.color_type = "OBJECT"
+
 
 class ColourByPset(bpy.types.Operator):
     bl_idname = "bim.colour_by_pset"
@@ -164,6 +186,14 @@ class ColourByPset(bpy.types.Operator):
     bl_options = {"REGISTER", "UNDO"}
 
     def execute(self, context):
+        IfcStore.begin_transaction(self)
+        self.store_state()
+        result = self._execute(context)
+        IfcStore.add_transaction_operation(self)
+        IfcStore.end_transaction(self)
+        return result
+
+    def _execute(self, context):
         self.file = IfcStore.get_file()
         colours = cycle(colour_list)
         values = {}
@@ -188,10 +218,24 @@ class ColourByPset(bpy.types.Operator):
             if value not in values:
                 values[value] = next(colours)
             obj.color = values[value]
-        area = next(area for area in context.screen.areas if area.type == "VIEW_3D")
-        area.spaces[0].shading.color_type = "OBJECT"
+        areas = [a for a in bpy.context.screen.areas if a.type == "VIEW_3D"]
+        if areas:
+            areas[0].spaces[0].shading.color_type = "OBJECT"
         return {"FINISHED"}
 
+    def store_state(self):
+        areas = [a for a in bpy.context.screen.areas if a.type == "VIEW_3D"]
+        if areas:
+            self.transaction_data = {"area": areas[0], "color_type": areas[0].spaces[0].shading.color_type}
+
+    def rollback(self, data):
+        if data:
+            data["area"].spaces[0].shading.color_type = data["color_type"]
+
+    def commit(self, data):
+        if data:
+            data["area"].spaces[0].shading.color_type = "OBJECT"
+
 
 class ColourByClass(bpy.types.Operator):
     bl_idname = "bim.colour_by_class"
@@ -199,6 +243,14 @@ class ColourByClass(bpy.types.Operator):
     bl_options = {"REGISTER", "UNDO"}
 
     def execute(self, context):
+        IfcStore.begin_transaction(self)
+        self.store_state()
+        result = self._execute(context)
+        IfcStore.add_transaction_operation(self)
+        IfcStore.end_transaction(self)
+        return result
+
+    def _execute(self, context):
         self.file = IfcStore.get_file()
         colours = cycle(colour_list)
         ifc_classes = {}
@@ -210,10 +262,24 @@ class ColourByClass(bpy.types.Operator):
             if ifc_class not in ifc_classes:
                 ifc_classes[ifc_class] = next(colours)
             obj.color = ifc_classes[ifc_class]
-        area = next(area for area in bpy.context.screen.areas if area.type == "VIEW_3D")
-        area.spaces[0].shading.color_type = "OBJECT"
+        areas = [a for a in bpy.context.screen.areas if a.type == "VIEW_3D"]
+        if areas:
+            areas[0].spaces[0].shading.color_type = "OBJECT"
         return {"FINISHED"}
 
+    def store_state(self):
+        areas = [a for a in bpy.context.screen.areas if a.type == "VIEW_3D"]
+        if areas:
+            self.transaction_data = {"area": areas[0], "color_type": areas[0].spaces[0].shading.color_type}
+
+    def rollback(self, data):
+        if data:
+            data["area"].spaces[0].shading.color_type = data["color_type"]
+
+    def commit(self, data):
+        if data:
+            data["area"].spaces[0].shading.color_type = "OBJECT"
+
 
 class ResetObjectColours(bpy.types.Operator):
     bl_idname = "bim.reset_object_colours"