From f43e666583ddae351c3aa9d047b4a509e62566c8 Mon Sep 17 00:00:00 2001
From: Bruno Postle <bruno@postle.net>
Date: Sun, 23 Mar 2025 12:07:41 +0000
Subject: [PATCH] git colourisation fixes

Fix bug where uncommitted changes were not fully colourised in the same
way as diffs. Also try and catch more changes, eg. highlight if a
Products Type has changed.
---
 src/bonsai/bonsai/core/ifcgit.py | 10 +++++++---
 src/bonsai/bonsai/tool/ifcgit.py | 30 +++++++++++++++++++++---------
 2 files changed, 28 insertions(+), 12 deletions(-)

diff --git a/src/bonsai/bonsai/core/ifcgit.py b/src/bonsai/bonsai/core/ifcgit.py
index 731e3751f2..404bd8df9b 100644
--- a/src/bonsai/bonsai/core/ifcgit.py
+++ b/src/bonsai/bonsai/core/ifcgit.py
@@ -99,15 +99,19 @@ def colourise_revision(ifcgit: tool.IfcGit) -> None:
     step_ids = ifcgit.get_revisions_step_ids()
     if not step_ids:
         return
-    modified_shape_object_step_ids = ifcgit.get_modified_shape_object_step_ids(step_ids)
-    final_step_ids = ifcgit.update_step_ids(step_ids, modified_shape_object_step_ids)
+    modified_step_ids = ifcgit.get_modified_step_ids(step_ids)
+    final_step_ids = ifcgit.update_step_ids(step_ids, modified_step_ids)
     ifcgit.colourise(final_step_ids)
 
 
 def colourise_uncommitted(ifcgit: tool.IfcGit, ifc: tool.Ifc, repo: git.Repo) -> None:
     path_ifc = ifc.get_path()
     step_ids = ifcgit.ifc_diff_ids(repo, None, "HEAD", path_ifc)
-    ifcgit.colourise(step_ids)
+    if not step_ids:
+        return
+    modified_step_ids = ifcgit.get_modified_step_ids(step_ids)
+    final_step_ids = ifcgit.update_step_ids(step_ids, modified_step_ids)
+    ifcgit.colourise(final_step_ids)
 
 
 def switch_revision(ifcgit: tool.IfcGit, ifc: tool.Ifc) -> None:
diff --git a/src/bonsai/bonsai/tool/ifcgit.py b/src/bonsai/bonsai/tool/ifcgit.py
index e306271ebd..1fe721d2f7 100644
--- a/src/bonsai/bonsai/tool/ifcgit.py
+++ b/src/bonsai/bonsai/tool/ifcgit.py
@@ -370,24 +370,36 @@ class IfcGit:
         return step_ids
 
     @classmethod
-    def get_modified_shape_object_step_ids(cls, step_ids: STEP_IDS) -> STEP_IDS:
+    def get_modified_step_ids(cls, step_ids: STEP_IDS) -> STEP_IDS:
         model = tool.Ifc.get()
-        modified_shape_object_step_ids = {"modified": []}
+        modified_step_ids = {"modified": set()}
 
-        for step_id in step_ids["modified"]:
-            if model.by_id(step_id).is_a() == "IfcProductDefinitionShape":
-                product = model.by_id(step_id).ShapeOfProduct[0]
-                modified_shape_object_step_ids["modified"].append(product.id())
+        for step_id in step_ids["modified"] | step_ids["added"]:
+            try:
+                entity = model.by_id(step_id)
+            except:
+                continue
+            if entity.is_a("IfcProductDefinitionShape"):
+                for product in entity.ShapeOfProduct:
+                    modified_step_ids["modified"].add(product.id())
+            elif entity.is_a("IfcObjectPlacement"):
+                for product in entity.PlacesObject:
+                    modified_step_ids["modified"].add(product.id())
+            elif entity.is_a("IfcTypeProduct") and entity.Types:
+                for related_object in entity.Types[0].RelatedObjects:
+                    modified_step_ids["modified"].add(related_object.id())
 
-        return modified_shape_object_step_ids
+        return modified_step_ids
 
     @classmethod
-    def update_step_ids(cls, step_ids: STEP_IDS, modified_shape_object_step_ids: STEP_IDS) -> STEP_IDS:
+    def update_step_ids(cls, step_ids: STEP_IDS, modified_step_ids: STEP_IDS) -> STEP_IDS:
 
         final_step_ids = {}
         final_step_ids["added"] = step_ids["added"]
         final_step_ids["removed"] = step_ids["removed"]
-        final_step_ids["modified"] = step_ids["modified"].union(modified_shape_object_step_ids["modified"])
+        final_step_ids["modified"] = (
+            step_ids["modified"].union(modified_step_ids["modified"]).difference(step_ids["added"])
+        )
         return final_step_ids
 
     @classmethod