From 60837b9512ec677093352b210c7d5e4418d41c0c Mon Sep 17 00:00:00 2001
From: Andrej730 <azhilenkov@gmail.com>
Date: Fri, 18 Oct 2024 17:08:37 +0500
Subject: [PATCH] Operator to save pset as a template

Mentioned in #5596

Example - https://imgchest.com/p/5xy235xvj4l

Adds 2 things:
1) UI indicator whether pset is based on template or not
2) When pset is not based on a template, indicator becomes clickable and works as a shortcut for creating pset template based on the current pset.
---
 src/bonsai/bonsai/bim/module/pset/__init__.py |  1 +
 src/bonsai/bonsai/bim/module/pset/data.py     |  2 ++
 src/bonsai/bonsai/bim/module/pset/operator.py | 31 +++++++++++++++++++
 src/bonsai/bonsai/bim/module/pset/ui.py       |  5 +++
 src/bonsai/bonsai/tool/pset_template.py       | 16 ++++++++++
 5 files changed, 55 insertions(+)

diff --git a/src/bonsai/bonsai/bim/module/pset/__init__.py b/src/bonsai/bonsai/bim/module/pset/__init__.py
index 14706e75be..5e53bde928 100644
--- a/src/bonsai/bonsai/bim/module/pset/__init__.py
+++ b/src/bonsai/bonsai/bim/module/pset/__init__.py
@@ -28,6 +28,7 @@ classes = (
     operator.EditPset,
     operator.EnablePsetEditing,
     operator.RemovePset,
+    operator.SavePsetAsTemplate,
     operator.TogglePsetExpansion,
     operator.UnsharePset,
     operator.BIM_OT_add_property_to_edit,
diff --git a/src/bonsai/bonsai/bim/module/pset/data.py b/src/bonsai/bonsai/bim/module/pset/data.py
index 4fe04538b7..7f1e391629 100644
--- a/src/bonsai/bonsai/bim/module/pset/data.py
+++ b/src/bonsai/bonsai/bim/module/pset/data.py
@@ -56,6 +56,7 @@ class Data:
         for name, data in sorted(psetqtos.items()):
             pset = ifc_file.by_id(data["id"])
             pset_uses = ifcopenshell.util.element.get_elements_by_pset(pset)
+            has_template = bool(tool.Pset.get_pset_template(name))
             results.append(
                 {
                     "id": data["id"],
@@ -63,6 +64,7 @@ class Data:
                     "is_expanded": is_expanded.get(data["id"], True),
                     "Properties": [{"Name": k, "NominalValue": v} for k, v in sorted(data.items()) if k != "id"],
                     "shared_pset_uses": len(pset_uses),
+                    "has_template": has_template,
                 }
             )
         return sorted(results, key=lambda v: v["Name"])
diff --git a/src/bonsai/bonsai/bim/module/pset/operator.py b/src/bonsai/bonsai/bim/module/pset/operator.py
index fe08585701..3d76f61c0b 100644
--- a/src/bonsai/bonsai/bim/module/pset/operator.py
+++ b/src/bonsai/bonsai/bim/module/pset/operator.py
@@ -453,3 +453,34 @@ class AddProposedProp(bpy.types.Operator):
             self.report({"ERROR"}, res)
             return {"CANCELLED"}
         return {"FINISHED"}
+
+
+class SavePsetAsTemplate(bpy.types.Operator, tool.PsetTemplate.PsetTemplateOperator):
+    bl_idname = "bim.save_pset_as_template"
+    bl_label = "Save Pset As Template"
+    bl_description = "Save the provided pset as a pset template"
+    bl_options = {"REGISTER", "UNDO"}
+    pset_id: bpy.props.IntProperty()
+
+    def invoke(self, context, event):
+        props = context.scene.BIMPsetTemplateProperties
+        if tool.Blender.get_enum_safe(props, "pset_template_files") is None:
+            self.report({"ERROR"}, "No template files found. You can create one in Property Set Templates UI.")
+            return {"CANCELLED"}
+        return context.window_manager.invoke_props_dialog(self, width=250)
+
+    def draw(self, context):
+        props = context.scene.BIMPsetTemplateProperties
+        self.layout.prop(props, "pset_template_files", text="Template File")
+
+    def _execute(self, context):
+        ifc_file = tool.Ifc.get()
+        pset = ifc_file.by_id(self.pset_id)
+        template_file = IfcStore.pset_template_file
+        assert template_file
+
+        tool.PsetTemplate.add_pset_as_template(pset, template_file)
+
+        template_file.write(IfcStore.pset_template_path)
+        bonsai.bim.handler.refresh_ui_data()
+        bonsai.bim.schema.reload(ifc_file.schema)
diff --git a/src/bonsai/bonsai/bim/module/pset/ui.py b/src/bonsai/bonsai/bim/module/pset/ui.py
index 1625bc8966..baf694255e 100644
--- a/src/bonsai/bonsai/bim/module/pset/ui.py
+++ b/src/bonsai/bonsai/bim/module/pset/ui.py
@@ -135,6 +135,11 @@ def draw_psetqto_ui(
         op.pset_id = pset_id
         op.obj = obj_name
         op.obj_type = obj_type
+        if pset["has_template"]:
+            row.label(text="", icon="ASSET_MANAGER")
+        else:
+            op = row.operator("bim.save_pset_as_template", icon="ASSET_MANAGER", text="")
+            op.pset_id = pset_id
         remove_pset_row = row.row(align=True)
         op = remove_pset_row.operator("bim.remove_pset", icon="X", text="")
         op.pset_id = pset_id
diff --git a/src/bonsai/bonsai/tool/pset_template.py b/src/bonsai/bonsai/tool/pset_template.py
index 0307c1ed5d..9abbb2074f 100644
--- a/src/bonsai/bonsai/tool/pset_template.py
+++ b/src/bonsai/bonsai/tool/pset_template.py
@@ -54,3 +54,19 @@ class PsetTemplate(bonsai.core.tool.PsetTemplate):
 
         def _execute(self, context: bpy.types.Context) -> None:
             tool.Ifc.Operator._execute(self, context)
+
+    @classmethod
+    def add_pset_as_template(
+        cls, pset: ifcopenshell.entity_instance, template_file: ifcopenshell.file
+    ) -> ifcopenshell.entity_instance:
+        # TODO: add tests.
+        pset_template = ifcopenshell.api.pset_template.add_pset_template(template_file, pset.Name)
+        for property in pset.HasProperties:
+            ifcopenshell.api.pset_template.add_prop_template(
+                template_file,
+                pset_template,
+                name=property.Name,
+                description=property.Description,
+                primary_measure_type=property.NominalValue.is_a(),
+            )
+        return pset_template