From bec61d6d2c0367146c898f7fa806e9f817980b01 Mon Sep 17 00:00:00 2001
From: Petru Conduraru <petru@bimvoice.com>
Date: Thu, 16 Jul 2026 13:40:20 +0300
Subject: [PATCH] Bonsai: persist colour-by-property Pset selection across
 save/reload

SaveColourscheme only serialized colourscheme_query (the free text
custom-query field), never colourscheme_key (the dropdown that
actually drives colour-by-property when not in Custom Query mode).
LoadColourscheme was symmetrically blind to it too, so after saving a
colour scheme built from a Pset property and reopening the file, the
colours came back but the Pset selection was silently lost, exactly
as reported.

Save now records whichever of colourscheme_key/colourscheme_query was
actually driving the query, matching the logic already used by
ColourByProperty and SelectByProperty. Load restores colourscheme_key,
falling back to "QUERY" when the saved key is no longer a valid choice
for the current active object (its dynamic enum is scoped per-object),
so the saved query string still gets used instead of the selection
silently disappearing.

Fixes #6354.

Generated with the assistance of an AI coding tool.
---
 .../bonsai/bim/module/search/operator.py      | 28 +++++++++++++++++--
 1 file changed, 26 insertions(+), 2 deletions(-)

diff --git a/src/bonsai/bonsai/bim/module/search/operator.py b/src/bonsai/bonsai/bim/module/search/operator.py
index f55dcf31b7..f4d447ac23 100644
--- a/src/bonsai/bonsai/bim/module/search/operator.py
+++ b/src/bonsai/bonsai/bim/module/search/operator.py
@@ -39,6 +39,7 @@ from natsort import natsorted
 import bonsai.core.search as core
 import bonsai.tool as tool
 from bonsai.bim.ifc import IfcStore
+from bonsai.bim.module.search.data import ColourByPropertyData
 from bonsai.bim.prop import StrProperty
 
 if TYPE_CHECKING:
@@ -1186,7 +1187,8 @@ class SaveColourscheme(Operator, tool.Ifc.Operator):
             return
 
         props = tool.Search.get_search_props()
-        query = props.colourscheme_query
+        colourscheme_key = props.colourscheme_key
+        query = props.colourscheme_query if colourscheme_key == "QUERY" else colourscheme_key
 
         group = [g for g in tool.Ifc.get().by_type("IfcGroup") if g.Name == self.name]
         colourscheme = {cs.name: {"colour": cs.colour[0:3], "total": cs.total} for cs in props.colourscheme}
@@ -1195,9 +1197,17 @@ class SaveColourscheme(Operator, tool.Ifc.Operator):
             description = json.loads(group.Description)
             description["colourscheme"] = colourscheme
             description["colourscheme_query"] = query
+            description["colourscheme_key"] = colourscheme_key
             group.Description = json.dumps(description)
         else:
-            description = json.dumps({"type": "BBIM_Search", "colourscheme": colourscheme, "colourscheme_query": query})
+            description = json.dumps(
+                {
+                    "type": "BBIM_Search",
+                    "colourscheme": colourscheme,
+                    "colourscheme_query": query,
+                    "colourscheme_key": colourscheme_key,
+                }
+            )
             group = ifcopenshell.api.group.add_group(tool.Ifc.get(), name=self.name, description=description)
 
     def invoke(self, context, event):
@@ -1215,6 +1225,20 @@ class LoadColourscheme(Operator, tool.Ifc.Operator):
         group = tool.Ifc.get().by_id(int(props.saved_colourschemes))
         description = json.loads(group.Description)
         props.colourscheme_query = description.get("colourscheme_query")
+
+        # colourscheme_key is a dynamic enum scoped to the active object's Psets
+        # (see ColourByPropertyData.colourscheme_key), so a previously saved key
+        # is only a valid choice if the current active object still exposes it.
+        # Otherwise fall back to "QUERY" so the saved query string (restored
+        # above) is used as a custom query instead of silently dropping the
+        # colour-by-property selection.
+        colourscheme_key = description.get("colourscheme_key")
+        valid_keys = {item[0] for item in ColourByPropertyData.colourscheme_key() if item is not None}
+        if colourscheme_key and colourscheme_key in valid_keys:
+            props.colourscheme_key = colourscheme_key
+        else:
+            props.colourscheme_key = "QUERY"
+
         props.colourscheme.clear()
         for name, data in description.get("colourscheme", {}).items():
             new = props.colourscheme.add()