typing

src/bonsai/bonsai/bim/helper.py

@@ -152,7 +152,7 @@ def import_attribute(
     if isinstance(data_type, tuple) or data_type == "entity":
         callback(attribute.name(), None, data) if callback else None
         return
-    new = props.add()
+    new: bonsai.bim.prop.Attribute = props.add()
     new.name = attribute.name()
     new.is_null = data[attribute.name()] is None
     new.is_optional = attribute.optional()

src/bonsai/bonsai/bim/module/library/prop.py

@@ -42,7 +42,16 @@ class LibraryReference(PropertyGroup):
 
 
 class BIMLibraryProperties(PropertyGroup):
-    editing_mode: StringProperty(name="Editing Mode")
+    editing_mode: EnumProperty(
+        name="Editing Mode",
+        items=(
+            ("-", "-", ""),
+            ("LIBRARY", "LIBRARY", ""),
+            ("REFERENCES", "REFERENCES", ""),
+            ("REFERENCE", "REFERENCE", ""),
+        ),
+        default="-",
+    )
     library_attributes: CollectionProperty(name="Library Attributes", type=Attribute)
     active_library_id: IntProperty(name="Active Library Id")
     reference_attributes: CollectionProperty(name="Library Attributes", type=Attribute)

src/bonsai/bonsai/bim/module/pset/prop.py

@@ -38,8 +38,7 @@ from bpy.props import (
     FloatVectorProperty,
     CollectionProperty,
 )
-from typing import Union, Literal
-
+from typing import Union, Literal, TYPE_CHECKING, get_args
 
 psetnames = {}
 qtonames = {}
@@ -231,13 +230,18 @@ class IfcPropertyEnumeratedValue(PropertyGroup):
     enumerated_values: CollectionProperty(type=Attribute)
 
 
+IfcPropertyValueType = Literal["IfcPropertySingleValue", "IfcPropertyEnumeratedValue"]
+
+
 class IfcProperty(PropertyGroup):
     metadata: PointerProperty(type=Attribute)
-    value_type: EnumProperty(
-        items=[(v, v, v) for v in ("IfcPropertySingleValue", "IfcPropertyEnumeratedValue")], name="Value Type"
-    )
+    value_type: EnumProperty(items=[(v, v, v) for v in get_args(IfcPropertyValueType)], name="Value Type")
     enumerated_value: PointerProperty(type=IfcPropertyEnumeratedValue)
 
+    if TYPE_CHECKING:
+        metadata: Attribute
+        value_type: IfcPropertyValueType
+
 
 class PsetProperties(PropertyGroup):
     active_pset_id: IntProperty(name="Active Pset ID")

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

@@ -16,6 +16,7 @@
 # You should have received a copy of the GNU General Public License
 # along with Bonsai.  If not, see <http://www.gnu.org/licenses/>.
 
+from __future__ import annotations
 import bpy
 import bonsai.tool as tool
 from bpy.types import Panel
@@ -35,21 +36,23 @@ from bonsai.bim.module.pset.data import (
     WorkSchedulePsetsData,
 )
 from bonsai.bim.module.material.data import ObjectMaterialData
-from typing import Any, Optional
+from typing import Any, Optional, TYPE_CHECKING
+from typing_extensions import assert_never
+
+if TYPE_CHECKING:
+    from bonsai.bim.module.pset.prop import IfcProperty
 
 
-def draw_property(
-    prop: bpy.types.PropertyGroup, layout: bpy.types.UILayout, copy_operator: Optional[str] = None
-) -> None:
+def draw_property(prop: IfcProperty, layout: bpy.types.UILayout, copy_operator: Optional[str] = None) -> None:
     if prop.value_type == "IfcPropertySingleValue":
         draw_single_property(prop, layout, copy_operator)
     elif prop.value_type == "IfcPropertyEnumeratedValue":
         draw_enumerated_property(prop, layout, copy_operator)
+    else:
+        assert_never(prop.value_type)
 
 
-def draw_single_property(
-    prop: bpy.types.PropertyGroup, layout: bpy.types.UILayout, copy_operator: Optional[str] = None
-) -> None:
+def draw_single_property(prop: IfcProperty, layout: bpy.types.UILayout, copy_operator: Optional[str] = None) -> None:
     value_name = prop.metadata.get_value_name(display_only=True)
     if not value_name:
         layout.label(text=prop["Name"])

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

@@ -43,8 +43,8 @@ class LoadGroupDecorationData:
         cls.is_loaded = True
 
     @classmethod
-    def load_groups_to_show(cls):
-        ret = []
+    def load_groups_to_show(cls) -> list[tuple[str, str, str]]:
+        ret: list[tuple[str, str, str]] = []
         abrv = {
             "LOAD_CASE": "L.Case: ",
             "LOAD_COMBINATION": "L.Comb: ",
@@ -64,7 +64,7 @@ class LoadGroupDecorationData:
                     for subgoup in [sg for sg in item if sg.is_a("IfcStructuralLoadGroup")]:
                         ret.append((str(subgoup.id()), ".       " + abrv[subgoup.PredefinedType] + subgoup.Name, ""))
 
-        if props.activity_type == "External Reaction":
+        elif props.activity_type == "External Reaction":
             groups = m.HasResults or []
             for g in groups:
                 result_name = g.ResultForLoadGroup.Name or ""

src/bonsai/bonsai/core/library.py

@@ -16,35 +16,43 @@
 # You should have received a copy of the GNU General Public License
 # along with Bonsai.  If not, see <http://www.gnu.org/licenses/>.
 
+from __future__ import annotations
+from typing import TYPE_CHECKING, Optional
 
-def add_library(ifc):
+if TYPE_CHECKING:
+    import bpy
+    import ifcopenshell
+    import bonsai.tool as tool
+
+
+def add_library(ifc: tool.Ifc) -> ifcopenshell.entity_instance:
     return ifc.run("library.add_library", name="Unnamed")
 
 
-def remove_library(ifc, library=None):
+def remove_library(ifc: tool.Ifc, library: ifcopenshell.entity_instance) -> None:
     ifc.run("library.remove_library", library=library)
 
 
-def enable_editing_library_references(library_tool, library=None):
+def enable_editing_library_references(library_tool: tool.Library, library: ifcopenshell.entity_instance) -> None:
     library_tool.set_editing_mode("REFERENCES")
     library_tool.set_active_library(library)
     library_tool.import_references(library)
 
 
-def disable_editing_library_references(library):
+def disable_editing_library_references(library: tool.Library) -> None:
     library.clear_editing_mode()
 
 
-def enable_editing_library(library):
+def enable_editing_library(library: tool.Library) -> None:
     library.set_editing_mode("LIBRARY")
     library.import_library_attributes(library.get_active_library())
 
 
-def disable_editing_library(library):
+def disable_editing_library(library: tool.Library) -> None:
     library.set_editing_mode("REFERENCES")
 
 
-def edit_library(ifc, library):
+def edit_library(ifc: tool.Ifc, library: tool.Library) -> None:
     library.set_editing_mode("REFERENCES")
     active_library = library.get_active_library()
     attributes = library.export_library_attributes()
@@ -52,29 +60,29 @@ def edit_library(ifc, library):
     library.import_references(active_library)
 
 
-def add_library_reference(ifc, library):
+def add_library_reference(ifc: tool.Ifc, library: tool.Library) -> ifcopenshell.entity_instance:
     active_library = library.get_active_library()
     reference = ifc.run("library.add_reference", library=active_library)
     library.import_references(active_library)
     return reference
 
 
-def remove_library_reference(ifc, library, reference=None):
+def remove_library_reference(ifc: tool.Ifc, library: tool.Library, reference: ifcopenshell.entity_instance) -> None:
     ifc.run("library.remove_reference", reference=reference)
     library.import_references(library.get_active_library())
 
 
-def enable_editing_library_reference(library, reference=None):
+def enable_editing_library_reference(library: tool.Library, reference: ifcopenshell.entity_instance) -> None:
     library.set_editing_mode("REFERENCE")
     library.set_active_reference(reference)
     library.import_reference_attributes(reference)
 
 
-def disable_editing_library_reference(library):
+def disable_editing_library_reference(library: tool.Library) -> None:
     library.set_editing_mode("REFERENCES")
 
 
-def edit_library_reference(ifc, library):
+def edit_library_reference(ifc: tool.Ifc, library: tool.Library) -> None:
     library.set_editing_mode("REFERENCES")
     active_reference = library.get_active_reference()
     attributes = library.export_reference_attributes()
@@ -82,9 +90,9 @@ def edit_library_reference(ifc, library):
     library.import_references(library.get_active_library())
 
 
-def assign_library_reference(ifc, obj=None, reference=None):
+def assign_library_reference(ifc: tool.Ifc, obj: bpy.types.Object, reference: ifcopenshell.entity_instance) -> None:
     ifc.run("library.assign_reference", products=[ifc.get_entity(obj)], reference=reference)
 
 
-def unassign_library_reference(ifc, obj=None, reference=None):
+def unassign_library_reference(ifc: tool.Ifc, obj: bpy.types.Object, reference: ifcopenshell.entity_instance) -> None:
     ifc.run("library.unassign_reference", products=[ifc.get_entity(obj)], reference=reference)

src/bonsai/bonsai/core/structural.py

@@ -16,35 +16,48 @@
 # You should have received a copy of the GNU General Public License
 # along with Bonsai.  If not, see <http://www.gnu.org/licenses/>.
 
+from __future__ import annotations
+from typing import TYPE_CHECKING, Optional, Union
 
-def add_structural_analysis_model(ifc, structural):
+if TYPE_CHECKING:
+    import bpy
+    import ifcopenshell
+    import bonsai.tool as tool
+
+
+def add_structural_analysis_model(ifc: tool.Ifc, structural: tool.Structural) -> ifcopenshell.entity_instance:
     result = ifc.run("structural.add_structural_analysis_model")
     structural.load_structural_analysis_models()
     structural.ensure_representation_contexts()
     return result
 
 
-def assign_structural_analysis_model(ifc, structural, product=None, structural_analysis_model=None):
-    product = structural.get_product_or_active_object(product)
+def assign_structural_analysis_model(
+    ifc: tool.Ifc,
+    structural: tool.Structural,
+    product: str,
+    structural_analysis_model: Union[int, None],
+) -> None:
+    product_ = structural.get_product_or_active_object(product)
     if structural_analysis_model and product:
         ifc.run(
             "structural.assign_structural_analysis_model",
             **{
-                "product": ifc.get().by_id(product.BIMObjectProperties.ifc_definition_id),
+                "product": ifc.get().by_id(product_.BIMObjectProperties.ifc_definition_id),
                 "structural_analysis_model": ifc.get().by_id(structural_analysis_model),
             },
         )
 
 
-def disable_editing_structural_analysis_model(structural):
+def disable_editing_structural_analysis_model(structural: tool.Structural) -> None:
     structural.disable_editing_structural_analysis_model()
 
 
-def disable_structural_analysis_model_editing_ui(structural):
+def disable_structural_analysis_model_editing_ui(structural: tool.Structural) -> None:
     structural.disable_structural_analysis_model_editing_ui()
 
 
-def edit_structural_analysis_model(ifc, structural):
+def edit_structural_analysis_model(ifc: tool.Ifc, structural: tool.Structural) -> None:
     attributes = structural.get_structural_analysis_model_attributes()
     ifc.run(
         "structural.edit_structural_analysis_model",
@@ -57,26 +70,28 @@ def edit_structural_analysis_model(ifc, structural):
     structural.disable_editing_structural_analysis_model()
 
 
-def enable_editing_structural_analysis_model(structural, model=None):
+def enable_editing_structural_analysis_model(structural: tool.Structural, model: Union[int, None]) -> None:
     structural.enable_editing_structural_analysis_model(model)
 
 
-def enable_structural_analysis_model_editing_ui(structural):
+def enable_structural_analysis_model_editing_ui(structural: tool.Structural) -> None:
     structural.enable_structural_analysis_model_editing_ui()
 
 
-def load_structural_analysis_model_attributes(structural, model=None):
+def load_structural_analysis_model_attributes(structural: tool.Structural, model: Union[int, None]) -> None:
     data = structural.get_ifc_structural_analysis_model_attributes(model)
+    if data is None:
+        return
     structural.load_structural_analysis_model_attributes(data)
 
 
-def load_structural_analysis_models(structural):
+def load_structural_analysis_models(structural: tool.Structural) -> None:
     structural.load_structural_analysis_models()
     structural.enable_structural_analysis_model_editing_ui()
     # structural.disable_editing_structural_analysis_model()
 
 
-def remove_structural_analysis_model(ifc, structural, model=None):
+def remove_structural_analysis_model(ifc: tool.Ifc, structural: tool.Structural, model: int) -> None:
     ifc.run(
         "structural.remove_structural_analysis_model",
         **{"structural_analysis_model": ifc.get().by_id(model)},
@@ -84,13 +99,18 @@ def remove_structural_analysis_model(ifc, structural, model=None):
     structural.load_structural_analysis_models()
 
 
-def unassign_structural_analysis_model(ifc, structural, product=None, structural_analysis_model=None):
-    product = structural.get_product_or_active_object(product)
+def unassign_structural_analysis_model(
+    ifc: tool.Ifc,
+    structural: tool.Structural,
+    product: str,
+    structural_analysis_model: Union[int, None],
+) -> None:
+    product_ = structural.get_product_or_active_object(product)
     if structural_analysis_model and product:
         ifc.run(
             "structural.unassign_structural_analysis_model",
             **{
-                "product": ifc.get().by_id(product.BIMObjectProperties.ifc_definition_id),
+                "product": ifc.get().by_id(product_.BIMObjectProperties.ifc_definition_id),
                 "structural_analysis_model": ifc.get().by_id(structural_analysis_model),
             },
         )

src/bonsai/bonsai/tool/library.py

@@ -16,48 +16,51 @@
 # You should have received a copy of the GNU General Public License
 # along with Bonsai.  If not, see <http://www.gnu.org/licenses/>.
 
+import ifcopenshell
 import bpy
+import bonsai.bim.helper
 import bonsai.core.tool
 import bonsai.tool as tool
+from typing import Literal, Any, Union
 
 
 class Library(bonsai.core.tool.Library):
     @classmethod
-    def clear_editing_mode(cls):
-        bpy.context.scene.BIMLibraryProperties.editing_mode = ""
+    def clear_editing_mode(cls) -> None:
+        bpy.context.scene.BIMLibraryProperties.editing_mode = "-"
 
     @classmethod
-    def export_library_attributes(cls):
+    def export_library_attributes(cls) -> dict[str, Any]:
         props = bpy.context.scene.BIMLibraryProperties
         return bonsai.bim.helper.export_attributes(props.library_attributes)
 
     @classmethod
-    def export_reference_attributes(cls):
+    def export_reference_attributes(cls) -> dict[str, Any]:
         props = bpy.context.scene.BIMLibraryProperties
         return bonsai.bim.helper.export_attributes(props.reference_attributes)
 
     @classmethod
-    def get_active_library(cls):
+    def get_active_library(cls) -> ifcopenshell.entity_instance:
         return tool.Ifc.get().by_id(bpy.context.scene.BIMLibraryProperties.active_library_id)
 
     @classmethod
-    def get_active_reference(cls):
+    def get_active_reference(cls) -> ifcopenshell.entity_instance:
         return tool.Ifc.get().by_id(bpy.context.scene.BIMLibraryProperties.active_reference_id)
 
     @classmethod
-    def import_library_attributes(cls, library):
+    def import_library_attributes(cls, library: ifcopenshell.entity_instance) -> None:
         props = bpy.context.scene.BIMLibraryProperties
         props.library_attributes.clear()
         bonsai.bim.helper.import_attributes2(library, props.library_attributes)
 
     @classmethod
-    def import_reference_attributes(cls, reference):
+    def import_reference_attributes(cls, reference: ifcopenshell.entity_instance) -> None:
         props = bpy.context.scene.BIMLibraryProperties
         props.reference_attributes.clear()
         bonsai.bim.helper.import_attributes2(reference, props.reference_attributes)
 
     @classmethod
-    def import_references(cls, library):
+    def import_references(cls, library: ifcopenshell.entity_instance) -> None:
         props = bpy.context.scene.BIMLibraryProperties
         props.references.clear()
         if tool.Ifc.get_schema() == "IFC2X3":
@@ -70,13 +73,13 @@ class Library(bonsai.core.tool.Library):
             new.name = reference.Name or "Unnamed"
 
     @classmethod
-    def set_active_library(cls, library):
+    def set_active_library(cls, library: ifcopenshell.entity_instance) -> None:
         bpy.context.scene.BIMLibraryProperties.active_library_id = library.id()
 
     @classmethod
-    def set_active_reference(cls, reference):
+    def set_active_reference(cls, reference: ifcopenshell.entity_instance) -> None:
         bpy.context.scene.BIMLibraryProperties.active_reference_id = reference.id()
 
     @classmethod
-    def set_editing_mode(cls, mode):
+    def set_editing_mode(cls, mode: Literal["LIBRARY", "REFERENCES", "REFERENCE"]) -> None:
         bpy.context.scene.BIMLibraryProperties.editing_mode = mode

src/bonsai/bonsai/tool/structural.py

@@ -18,38 +18,40 @@
 
 import bpy
 import ifcopenshell
+import ifcopenshell.util.representation
 import json
 import bonsai.bim.helper
 import bonsai.core.tool
 import bonsai.tool as tool
 from bonsai.bim.ifc import IfcStore
 from pprint import pprint
+from typing import Union, Any
 
 
 class Structural(bonsai.core.tool.Structural):
     @classmethod
-    def disable_editing_structural_analysis_model(cls):
+    def disable_editing_structural_analysis_model(cls) -> None:
         bpy.context.scene.BIMStructuralProperties.active_structural_analysis_model_id = 0
 
     @classmethod
-    def disable_structural_analysis_model_editing_ui(cls):
+    def disable_structural_analysis_model_editing_ui(cls) -> None:
         bpy.context.scene.BIMStructuralProperties.is_editing = False
 
     @classmethod
-    def enable_editing_structural_analysis_model(cls, model):
+    def enable_editing_structural_analysis_model(cls, model: Union[int, None]) -> None:
         if model:
             bpy.context.scene.BIMStructuralProperties.active_structural_analysis_model_id = model
 
     @classmethod
-    def enable_structural_analysis_model_editing_ui(cls):
+    def enable_structural_analysis_model_editing_ui(cls) -> None:
         bpy.context.scene.BIMStructuralProperties.is_editing = True
 
     @classmethod
-    def enabled_structural_analysis_model_editing_ui(cls):
+    def enabled_structural_analysis_model_editing_ui(cls) -> bool:
         return bpy.context.scene.BIMStructuralProperties.is_editing
 
     @classmethod
-    def ensure_representation_contexts(cls):
+    def ensure_representation_contexts(cls) -> None:
         model = ifcopenshell.util.representation.get_context(tool.Ifc.get(), "Model")
         if not model:
             model = tool.Ifc.run(
@@ -70,13 +72,13 @@ class Structural(bonsai.core.tool.Structural):
             )
 
     @classmethod
-    def get_active_structural_analysis_model(cls):
+    def get_active_structural_analysis_model(cls) -> ifcopenshell.entity_instance:
         props = bpy.context.scene.BIMStructuralProperties
         model = tool.Ifc.get().by_id(props.active_structural_analysis_model_id)
         return model
 
     @classmethod
-    def get_ifc_structural_analysis_model_attributes(cls, model):
+    def get_ifc_structural_analysis_model_attributes(cls, model: Union[int, None]) -> Union[dict[str, Any], None]:
         if model:
             ifc_model = tool.Ifc.get().by_id(model)
             data = ifc_model.get_info()
@@ -100,14 +102,14 @@ class Structural(bonsai.core.tool.Structural):
             return data
 
     @classmethod
-    def get_ifc_structural_analysis_models(cls):
+    def get_ifc_structural_analysis_models(cls) -> dict[int, dict[str, Union[str, None]]]:
         models = {}
         for model in tool.Ifc.get().by_type("IfcStructuralAnalysisModel"):
             models[model.id()] = {"Name": model.Name}
         return models
 
     @classmethod
-    def get_product_or_active_object(cls, product):
+    def get_product_or_active_object(cls, product: str) -> Union[bpy.types.Object, None]:
         product = bpy.data.objects.get(product) if product else bpy.context.active_object
         try:
             if product.BIMObjectProperties.ifc_definition_id:
@@ -118,13 +120,13 @@ class Structural(bonsai.core.tool.Structural):
             return None
 
     @classmethod
-    def get_structural_analysis_model_attributes(cls):
+    def get_structural_analysis_model_attributes(cls) -> dict[str, Any]:
         props = bpy.context.scene.BIMStructuralProperties
         attributes = bonsai.bim.helper.export_attributes(props.structural_analysis_model_attributes)
         return attributes
 
     @classmethod
-    def load_structural_analysis_model_attributes(cls, data):
+    def load_structural_analysis_model_attributes(cls, data: dict[str, Any]) -> None:
         props = bpy.context.scene.BIMStructuralProperties
         props.structural_analysis_model_attributes.clear()
         for attribute in IfcStore.get_schema().declaration_by_name("IfcStructuralAnalysisModel").all_attributes():
@@ -145,7 +147,7 @@ class Structural(bonsai.core.tool.Structural):
                 new.data_type = "string"
 
     @classmethod
-    def load_structural_analysis_models(cls):
+    def load_structural_analysis_models(cls) -> None:
         models = tool.Structural.get_ifc_structural_analysis_models()
         props = bpy.context.scene.BIMStructuralProperties
         props.structural_analysis_models.clear()

src/bonsai/test/tool/test_library.py

@@ -130,5 +130,5 @@ class TestSetActiveReference(NewFile):
 
 class TestSetEditingMode(NewFile):
     def test_run(self):
-        subject.set_editing_mode("FOO")
-        assert bpy.context.scene.BIMLibraryProperties.editing_mode == "FOO"
+        subject.set_editing_mode("LIBRARY")
+        assert bpy.context.scene.BIMLibraryProperties.editing_mode == "LIBRARY"