diff --git a/src/bonsai/bonsai/bim/module/classification/data.py b/src/bonsai/bonsai/bim/module/classification/data.py index b822903431..3ed42cf3ea 100644 --- a/src/bonsai/bonsai/bim/module/classification/data.py +++ b/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] = {} - is_loaded = False + obj_type = "Object" @classmethod - def load(cls): - 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"] = "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 + def get_element(cls) -> Union[ifcopenshell.entity_instance, None]: + if obj := bpy.context.active_object: + return tool.Ifc.get_entity(obj) + return None class MaterialClassificationsData(ReferencesData): - data: dict[str, Any] = {} - is_loaded = False + obj_type = "Material" @classmethod - def load(cls): - 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 = [] - + def get_element(cls) -> Union[ifcopenshell.entity_instance, None]: props = tool.Material.get_material_props() if material := props.active_material: - if material.ifc_definition_id: - element = tool.Ifc.get().by_id(material.ifc_definition_id) - for reference in ifcopenshell.util.classification.get_references(element): - data = reference.get_info() - del data["ReferencedSource"] - results.append(data) - return results + return tool.Ifc.get().by_id(material.ifc_definition_id) + return None class CostClassificationsData(ReferencesData): - data: dict[str, Any] = {} - is_loaded = False + obj_type = "Cost" @classmethod - def load(cls): - 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]] = [] + def get_element(cls) -> Union[ifcopenshell.entity_instance, None]: props = tool.Cost.get_cost_props() - active_cost_item = props.active_cost_item - if not active_cost_item: - return results - 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 + if cost_item := props.active_cost_item: + return tool.Ifc.get().by_id(cost_item.ifc_definition_id) + return None diff --git a/src/bonsai/bonsai/bim/module/classification/ui.py b/src/bonsai/bonsai/bim/module/classification/ui.py index 18038b539e..ddf5d5f2e1 100644 --- a/src/bonsai/bonsai/bim/module/classification/ui.py +++ b/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 False - return bool(tool.Ifc.get_entity(context.active_object)) + return bool((obj := context.active_object) and tool.Ifc.get_entity(obj)) - def draw(self, context): - if not ObjectClassificationsData.is_loaded: - ObjectClassificationsData.load() - self.data = ObjectClassificationsData - self.obj = context.active_object.name - self.obj_type = "Object" - self.draw_ui(context) + def get_object_name(self, context: bpy.types.Context) -> str: + assert (obj := context.active_object) + return obj.name 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(