ClassificationsData - simplify

src/bonsai/bonsai/bim/module/classification/data.py

@@ -21,7 +21,7 @@ import ifcopenshell
 import ifcopenshell.util.date
 import ifcopenshell.util.classification
 import bonsai.tool as tool
-from typing import Any
+from typing import Any, Literal, Union
 from bonsai.bim.ifc import IfcStore
 
 
@@ -78,6 +78,34 @@ class ClassificationsData:
 
 
 class ReferencesData:
+    data: dict[str, Any]
+    is_loaded = False
+    obj_type: Literal["Object", "Material", "Cost"]
+
+    @classmethod
+    def load(cls):
+        cls.is_loaded = True
+        cls.data = {}
+        cls.data["references"] = cls.references()
+        cls.data["active_classification_library"] = cls.active_classification_library()
+        cls.data["classifications"] = cls.classifications()
+
+    @classmethod
+    def get_element(cls) -> Union[ifcopenshell.entity_instance, None]:
+        """Get element to get classification references from."""
+        raise NotImplementedError
+
+    @classmethod
+    def references(cls) -> list[dict[str, Any]]:
+        results = []
+        element = cls.get_element()
+        if element:
+            for reference in ifcopenshell.util.classification.get_references(element):
+                data = reference.get_info()
+                del data["ReferencedSource"]
+                results.append(data)
+        return results
+
     @classmethod
     def active_classification_library(cls):
         if not IfcStore.classification_file or not IfcStore.classification_file.by_type("IfcClassification"):
@@ -93,79 +121,32 @@ class ReferencesData:
 
 
 class ObjectClassificationsData(ReferencesData):
-    data: dict[str, Any] = {}
+    obj_type = "Object"
-    is_loaded = False
 
     @classmethod
-    def load(cls):
+    def get_element(cls) -> Union[ifcopenshell.entity_instance, None]:
-        cls.is_loaded = True
+        if obj := bpy.context.active_object:
-        cls.data["references"] = cls.references()
+            return tool.Ifc.get_entity(obj)
-        cls.data["active_classification_library"] = cls.active_classification_library()
+        return None
-        cls.data["classifications"] = cls.classifications()
-        cls.data["object_type"] = "Object"
-
-    @classmethod
-    def references(cls):
-        results = []
-        element = tool.Ifc.get_entity(bpy.context.active_object)
-        if element:
-            for reference in ifcopenshell.util.classification.get_references(element):
-                data = reference.get_info()
-                del data["ReferencedSource"]
-                results.append(data)
-        return results
 
 
 class MaterialClassificationsData(ReferencesData):
-    data: dict[str, Any] = {}
+    obj_type = "Material"
-    is_loaded = False
 
     @classmethod
-    def load(cls):
+    def get_element(cls) -> Union[ifcopenshell.entity_instance, None]:
-        cls.is_loaded = True
-        cls.data["references"] = cls.references()
-        cls.data["active_classification_library"] = cls.active_classification_library()
-        cls.data["classifications"] = cls.classifications()
-        cls.data["object_type"] = "Material"
-
-    @classmethod
-    def references(cls):
-        results = []
-
         props = tool.Material.get_material_props()
         if material := props.active_material:
-            if material.ifc_definition_id:
+            return tool.Ifc.get().by_id(material.ifc_definition_id)
-                element = tool.Ifc.get().by_id(material.ifc_definition_id)
+        return None
-                for reference in ifcopenshell.util.classification.get_references(element):
-                    data = reference.get_info()
-                    del data["ReferencedSource"]
-                    results.append(data)
-        return results
 
 
 class CostClassificationsData(ReferencesData):
-    data: dict[str, Any] = {}
+    obj_type = "Cost"
-    is_loaded = False
 
     @classmethod
-    def load(cls):
+    def get_element(cls) -> Union[ifcopenshell.entity_instance, None]:
-        cls.is_loaded = True
-        cls.data["references"] = cls.references()
-        cls.data["active_classification_library"] = cls.active_classification_library()
-        cls.data["classifications"] = cls.classifications()
-        cls.data["object_type"] = "Cost"
-
-    @classmethod
-    def references(cls) -> list[dict[str, Any]]:
-        results: list[dict[str, Any]] = []
         props = tool.Cost.get_cost_props()
-        active_cost_item = props.active_cost_item
+        if cost_item := props.active_cost_item:
-        if not active_cost_item:
+            return tool.Ifc.get().by_id(cost_item.ifc_definition_id)
-            return results
+        return None
-        element = tool.Ifc.get().by_id(active_cost_item.ifc_definition_id)
-        if element:
-            for reference in ifcopenshell.util.classification.get_references(element):
-                data = reference.get_info()
-                del data["ReferencedSource"]
-                results.append(data)
-        return results

src/bonsai/bonsai/bim/module/classification/ui.py

@@ -112,7 +112,7 @@ class BIM_PT_classifications(Panel):
 
 
 class ReferenceUI:
-    layout: bpy.types.UILayout
+    layout: Union[bpy.types.UILayout, None]
     data: type[
         Union[
             ObjectClassificationsData,
@@ -121,7 +121,19 @@ class ReferenceUI:
         ]
     ]
 
-    def draw_ui(self, context) -> None:
+    obj: str
+    """Object name."""
+
+    def get_object_name(self, context: bpy.types.Context) -> str:
+        return ""
+
+    def draw(self, context: bpy.types.Context) -> None:
+        if not self.data.is_loaded:
+            self.data.load()
+        self.obj = self.get_object_name(context)
+        self.draw_ui(context)
+
+    def draw_ui(self, context: bpy.types.Context) -> None:
         self.sprops = tool.Classification.get_classification_props()
         self.bprops = tool.Bsdd.get_bsdd_props()
         self.props = tool.Classification.get_classification_reference_props()
@@ -139,7 +151,7 @@ class ReferenceUI:
             else:
                 self.draw_reference_ui(reference)
 
-    def draw_add_ui(self, context) -> None:
+    def draw_add_ui(self, context: bpy.types.Context) -> None:
         row = self.layout.row(align=True)
         row.label(text="Source", icon="OUTLINER")
         row.prop(self.sprops, "classification_source", text="")
@@ -151,20 +163,20 @@ class ReferenceUI:
         else:
             self.draw_add_bsdd_ui(context)
 
-    def draw_add_manual_ui(self, context) -> None:
+    def draw_add_manual_ui(self, context: object) -> None:
         row = self.layout.row()
         row.prop(self.props, "classifications", text="")
         if self.props.is_adding:
             bonsai.bim.helper.draw_attributes(self.props.reference_attributes, self.layout)
             row = self.layout.row(align=True)
             op = row.operator("bim.add_manual_classification_reference", text="Save", icon="CHECKMARK")
-            op.obj_type = self.data.data["object_type"]
+            op.obj_type = self.data.obj_type
             row.operator("bim.disable_adding_manual_classification_reference", text="", icon="CANCEL")
         else:
             row = self.layout.row()
             row.operator("bim.enable_adding_manual_classification_reference", text="Add Reference", icon="ADD")
 
-    def draw_add_bsdd_ui(self, context) -> None:
+    def draw_add_bsdd_ui(self, context: object) -> None:
         row = self.layout.row(align=True)
         row.prop(self.bprops, "keyword", text="")
         row.prop(self.bprops, "should_filter_ifc_class", text="", icon="FILTER")
@@ -189,9 +201,9 @@ class ReferenceUI:
                 "bim.add_classification_reference_from_bsdd", text="Add Classification Reference", icon="ADD"
             )
             op.obj = self.obj
-            op.obj_type = self.obj_type
+            op.obj_type = self.data.obj_type
 
-    def draw_add_file_ui(self, context) -> None:
+    def draw_add_file_ui(self, context: object) -> None:
         if not self.data.data["active_classification_library"]:
             row = self.layout.row(align=True)
             row.label(text="No Active Classification Library", icon="ERROR")
@@ -211,7 +223,7 @@ class ReferenceUI:
         if self.sprops.active_library_reference_index < len(self.sprops.available_library_references):
             op = row.operator("bim.add_classification_reference", text="", icon="ADD")
             op.obj = self.obj
-            op.obj_type = self.obj_type
+            op.obj_type = self.data.obj_type
             op.reference = self.sprops.available_library_references[
                 self.sprops.active_library_reference_index
             ].ifc_definition_id
@@ -247,7 +259,7 @@ class ReferenceUI:
         op = row.operator("bim.remove_classification_reference", text="", icon="X")
         op.reference = reference["id"]
         op.obj = self.obj
-        op.obj_type = self.obj_type
+        op.obj_type = self.data.obj_type
 
 
 class BIM_PT_classification_references(Panel, ReferenceUI):
@@ -259,19 +271,15 @@ class BIM_PT_classification_references(Panel, ReferenceUI):
     bl_context = "object"
     bl_parent_id = "BIM_PT_tab_object_metadata"
 
+    data = ObjectClassificationsData
+
     @classmethod
     def poll(cls, context):
-        if not context.active_object:
+        return bool((obj := context.active_object) and tool.Ifc.get_entity(obj))
-            return False
-        return bool(tool.Ifc.get_entity(context.active_object))
 
-    def draw(self, context):
+    def get_object_name(self, context: bpy.types.Context) -> str:
-        if not ObjectClassificationsData.is_loaded:
+        assert (obj := context.active_object)
-            ObjectClassificationsData.load()
+        return obj.name
-        self.data = ObjectClassificationsData
-        self.obj = context.active_object.name
-        self.obj_type = "Object"
-        self.draw_ui(context)
 
 
 class BIM_PT_material_classifications(Panel, ReferenceUI):
@@ -283,6 +291,8 @@ class BIM_PT_material_classifications(Panel, ReferenceUI):
     bl_context = "scene"
     bl_parent_id = "BIM_PT_materials"
 
+    data = MaterialClassificationsData
+
     @classmethod
     def poll(cls, context):
         if not tool.Ifc.get():
@@ -292,14 +302,6 @@ class BIM_PT_material_classifications(Panel, ReferenceUI):
             return True
         return False
 
-    def draw(self, context):
-        if not MaterialClassificationsData.is_loaded:
-            MaterialClassificationsData.load()
-        self.data = MaterialClassificationsData
-        self.obj = ""
-        self.obj_type = "Material"
-        self.draw_ui(context)
-
 
 class BIM_PT_cost_classifications(Panel, ReferenceUI):
     bl_label = "Cost Item Classifications"
@@ -310,6 +312,8 @@ class BIM_PT_cost_classifications(Panel, ReferenceUI):
     bl_context = "scene"
     bl_parent_id = "BIM_PT_cost_schedules"
 
+    data = CostClassificationsData
+
     @classmethod
     def poll(cls, context):
         if not tool.Ifc.get():
@@ -317,14 +321,6 @@ class BIM_PT_cost_classifications(Panel, ReferenceUI):
         props = tool.Cost.get_cost_props()
         return bool(props.cost_items)
 
-    def draw(self, context):
-        if not CostClassificationsData.is_loaded:
-            CostClassificationsData.load()
-        self.data = CostClassificationsData
-        self.obj = ""
-        self.obj_type = "Cost"
-        self.draw_ui(context)
-
 
 class BIM_UL_classifications(UIList):
     def draw_item(