From 5bd24b27603ecd19baade6bc6bd753786ced769c Mon Sep 17 00:00:00 2001
From: Andrej730 <azhilenkov@gmail.com>
Date: Mon, 27 Jan 2025 12:40:53 +0500
Subject: [PATCH] Enable pset template prop editing on adding a new one #6019

---
 .../bim/module/pset_template/operator.py      | 22 +++--------------
 .../bonsai/bim/module/pset_template/prop.py   |  8 +++++++
 src/bonsai/bonsai/tool/pset_template.py       | 24 +++++++++++++++++++
 3 files changed, 35 insertions(+), 19 deletions(-)

diff --git a/src/bonsai/bonsai/bim/module/pset_template/operator.py b/src/bonsai/bonsai/bim/module/pset_template/operator.py
index e9af89631e..251ebf6011 100644
--- a/src/bonsai/bonsai/bim/module/pset_template/operator.py
+++ b/src/bonsai/bonsai/bim/module/pset_template/operator.py
@@ -119,24 +119,7 @@ class EnableEditingPropTemplate(bpy.types.Operator):
     prop_template: bpy.props.IntProperty()
 
     def execute(self, context):
-        props = context.scene.BIMPsetTemplateProperties
-        props.active_prop_template_id = self.prop_template
-        template = IfcStore.pset_template_file.by_id(props.active_prop_template_id)
-        props.active_prop_template.name = template.Name or ""
-        props.active_prop_template.description = template.Description or ""
-        props.active_prop_template.primary_measure_type = template.PrimaryMeasureType or "-"
-        props.active_prop_template.template_type = template.TemplateType
-        props.active_prop_template.enum_values.clear()
-
-        if template.Enumerators:
-            props.active_prop_template.enum_values.clear()
-            data_type = props.active_prop_template.get_value_name()
-            for e in template.Enumerators.EnumerationValues:
-                new = props.active_prop_template.enum_values.add()
-                setattr(new, data_type, e.wrappedValue)
-
-        # Disable because of the intersecting enums in data.py.
-        props.active_pset_template_id = 0
+        tool.PsetTemplate.enable_editing_prop_template(tool.Ifc.get().by_id(self.prop_template))
         return {"FINISHED"}
 
 
@@ -245,7 +228,7 @@ class AddPropTemplate(bpy.types.Operator, tool.PsetTemplate.PsetTemplateOperator
     def _execute(self, context):
         props = context.scene.BIMPsetTemplateProperties
         pset_template_id = props.active_pset_template_id or int(props.pset_templates)
-        ifcopenshell.api.run(
+        prop_template = ifcopenshell.api.run(
             "pset_template.add_prop_template",
             IfcStore.pset_template_file,
             pset_template=IfcStore.pset_template_file.by_id(pset_template_id),
@@ -254,6 +237,7 @@ class AddPropTemplate(bpy.types.Operator, tool.PsetTemplate.PsetTemplateOperator
         IfcStore.pset_template_file.write(IfcStore.pset_template_path)
         bonsai.bim.handler.refresh_ui_data()
         bonsai.bim.schema.reload(tool.Ifc.get().schema)
+        tool.PsetTemplate.enable_editing_prop_template(prop_template)
 
 
 class RemovePropTemplate(bpy.types.Operator, tool.PsetTemplate.PsetTemplateOperator):
diff --git a/src/bonsai/bonsai/bim/module/pset_template/prop.py b/src/bonsai/bonsai/bim/module/pset_template/prop.py
index 0a0f707d72..66cd7e3fea 100644
--- a/src/bonsai/bonsai/bim/module/pset_template/prop.py
+++ b/src/bonsai/bonsai/bim/module/pset_template/prop.py
@@ -196,6 +196,14 @@ class PropTemplate(PropertyGroup):
         else:
             assert False, "Unknown data type"
 
+    if TYPE_CHECKING:
+        global_id: str
+        name: str
+        description: str
+        primary_measure_type: str
+        template_type: str
+        enum_values: bpy.types.bpy_prop_collection_idprop[EnumerationValues]
+
 
 class BIMPsetTemplateProperties(PropertyGroup):
     pset_template_files: EnumProperty(
diff --git a/src/bonsai/bonsai/tool/pset_template.py b/src/bonsai/bonsai/tool/pset_template.py
index bca0535290..45de43f3ac 100644
--- a/src/bonsai/bonsai/tool/pset_template.py
+++ b/src/bonsai/bonsai/tool/pset_template.py
@@ -95,3 +95,27 @@ class PsetTemplate(bonsai.core.tool.PsetTemplate):
 
         # Disable because of the intersecting enums in data.py.
         props.active_prop_template_id = 0
+
+    @classmethod
+    def enable_editing_prop_template(cls, prop_template: ifcopenshell.entity_instance) -> None:
+        props = tool.PsetTemplate.get_pset_template_props()
+        props.active_prop_template_id = prop_template.id()
+
+        pset_template_file = IfcStore.pset_template_file
+        assert pset_template_file
+        template = pset_template_file.by_id(props.active_prop_template_id)
+        props.active_prop_template.name = template.Name or ""
+        props.active_prop_template.description = template.Description or ""
+        props.active_prop_template.primary_measure_type = template.PrimaryMeasureType or "-"
+        props.active_prop_template.template_type = template.TemplateType
+        props.active_prop_template.enum_values.clear()
+
+        if template.Enumerators:
+            props.active_prop_template.enum_values.clear()
+            data_type = props.active_prop_template.get_value_name()
+            for e in template.Enumerators.EnumerationValues:
+                new = props.active_prop_template.enum_values.add()
+                setattr(new, data_type, e.wrappedValue)
+
+        # Disable because of the intersecting enums in data.py.
+        props.active_pset_template_id = 0