typing

src/bonsai/bonsai/bim/module/geometry/operator.py

@@ -23,6 +23,7 @@ import numpy as np
 import numpy.typing as npt
 import ifcopenshell
 import ifcopenshell.api.layer
+import ifcopenshell.api.root
 import ifcopenshell.api.style
 import ifcopenshell.util.element
 import ifcopenshell.util.placement
@@ -48,6 +49,9 @@ from ifcopenshell.util.shape_builder import ShapeBuilder
 from typing import Any, Union, Literal, get_args, TYPE_CHECKING, assert_never
 from bonsai.bim.module.model.decorator import ProfileDecorator
 
+if TYPE_CHECKING:
+    from bpy._typing import rna_enums
+
 
 class EditObjectPlacement(bpy.types.Operator, tool.Ifc.Operator):
     bl_idname = "bim.edit_object_placement"
@@ -1006,10 +1010,12 @@ class OverrideDuplicateMove(bpy.types.Operator):
         return OverrideDuplicateMove.execute_ifc_duplicate_operator(self, context)
 
     @staticmethod
-    def execute_duplicate_operator(self, context, linked=False):
+    def execute_duplicate_operator(
+        operator: bpy.types.Operator, context: bpy.types.Context, linked: bool = False
+    ) -> set["rna_enums.OperatorReturnItems"]:
         # Deep magick from the dawn of time
         if tool.Ifc.get():
-            IfcStore.execute_ifc_operator(self, context)
+            IfcStore.execute_ifc_operator(operator, context)
             return {"FINISHED"}
 
         if linked:
@@ -1019,15 +1025,15 @@ class OverrideDuplicateMove(bpy.types.Operator):
         return {"FINISHED"}
 
     @staticmethod
-    def execute_ifc_duplicate_operator(self, context, linked=False):
+    def execute_ifc_duplicate_operator(operator: bpy.types.Operator, context: bpy.types.Context, linked: bool = False):
         for obj in context.selected_objects:
             if element := tool.Ifc.get_entity(obj):
                 if element.is_a("IfcAnnotation") and element.ObjectType == "DRAWING":
                     tool.Blender.deselect_object(obj)
-                    self.report({"ERROR"}, "Drawing not duplicated.")
+                    operator.report({"ERROR"}, "Drawing not duplicated.")
                 elif tool.Geometry.is_locked(element):
                     tool.Blender.deselect_object(obj)
-                    self.report({"ERROR"}, lock_error_message(obj.name))
+                    operator.report({"ERROR"}, lock_error_message(obj.name))
         old_to_new, new_active_obj = tool.Geometry.duplicate_ifc_objects(
             set(context.selected_objects), linked=linked, active_object=context.active_object
         )
@@ -1557,7 +1563,7 @@ class OverrideJoin(bpy.types.Operator, tool.Ifc.Operator):
                     continue
                 element = tool.Ifc.get_entity(obj)
                 if element:
-                    ifcopenshell.api.run("root.remove_product", ifc_file, product=element)
+                    ifcopenshell.api.root.remove_product(ifc_file, product=element)
             bpy.ops.object.join()
             bpy.ops.bim.update_representation(obj=self.target.name, ifc_representation_class="")
         else:
@@ -1676,7 +1682,7 @@ class OverrideJoin(bpy.types.Operator, tool.Ifc.Operator):
                         assert False, f"Unexpected item type: {item.is_a()}. This is a bug."
 
                     items.append(copied_item)
-                ifcopenshell.api.run("root.remove_product", ifc_file, product=element)
+                ifcopenshell.api.root.remove_product(ifc_file, product=element)
             representation.Items = items
             bpy.ops.object.join()
             core.switch_representation(

src/bonsai/bonsai/core/drawing.py

@@ -312,7 +312,7 @@ def duplicate_drawing(
     ifc.run("group.edit_group", group=new_group, attributes={"Name": drawing_name, "ObjectType": "DRAWING"})
     ifc.run("group.assign_group", group=new_group, products=[new_drawing])
     if should_duplicate_annotations:
-        new_annotations = []
+        new_annotations: list[ifcopenshell.entity_instance] = []
         annotation_objs = [ifc.get_object(a) for a in drawing_tool.get_group_elements(group) if a != drawing]
         old_to_new, _ = geometry.duplicate_ifc_objects(annotation_objs)
         for new_elements in old_to_new.values():

src/bonsai/bonsai/tool/geometry.py

@@ -43,6 +43,7 @@ import ifcopenshell.util.unit
 import bonsai.core.tool
 import bonsai.core.drawing
 import bonsai.core.geometry
+import bonsai.core.root
 import bonsai.core.spatial
 import bonsai.core.style
 import bonsai.core.system
@@ -2013,8 +2014,11 @@ class Geometry(bonsai.core.tool.Geometry):
 
     @classmethod
     def duplicate_ifc_objects(
-        cls, objects_to_duplicate: Iterable[bpy.types.Object], active_object=None, linked=False
+        cls,
-    ) -> dict:
+        objects_to_duplicate: Iterable[bpy.types.Object],
+        active_object: Optional[bpy.types.Object] = None,
+        linked: bool = False,
+    ) -> tuple[dict[ifcopenshell.entity_instance, list[ifcopenshell.entity_instance]], Union[bpy.types.Object, None]]:
         # Handle arrays
         objects_to_duplicate = set(objects_to_duplicate)
         arrays_to_duplicate, array_children = cls.process_arrays_for_duplication(objects_to_duplicate)
@@ -2026,7 +2030,7 @@ class Geometry(bonsai.core.tool.Geometry):
         # Track decompositions so they can be recreated after the operation
         decomposition_relationships = tool.Root.get_decomposition_relationships(objects_to_duplicate)
         connection_relationships = tool.Root.get_connection_relationships(objects_to_duplicate)
-        old_to_new = {}
+        old_to_new: dict[ifcopenshell.entity_instance, list[ifcopenshell.entity_instance]] = {}
 
         for obj in objects_to_duplicate:
             element = tool.Ifc.get_entity(obj)

src/bonsai/bonsai/tool/ifc.py

@@ -114,7 +114,9 @@ class Ifc(bonsai.core.tool.Ifc):
         return IfcStore.get_schema()
 
     @classmethod
-    def get_entity(cls, obj: IFC_CONNECTED_TYPE) -> Union[ifcopenshell.entity_instance, None]:
+    def get_entity(
+        cls, obj: Union[IFC_CONNECTED_TYPE, tool.Geometry.TYPES_WITH_MESH_PROPERTIES]
+    ) -> Union[ifcopenshell.entity_instance, None]:
         """Get linked IFC entity based on obj's ifc_definition_id.
 
         Return None if object is not linked to IFC or it's linked to non-existent element.

src/bonsai/test/core/bootstrap.py

@@ -20,7 +20,7 @@ import sys
 import json
 import pytest
 import bonsai.core.tool
-from typing import Any, Optional
+from typing import Any, Optional, Type, Union, TypedDict, Literal
 from typing_extensions import Self
 
 
@@ -241,20 +241,43 @@ def voider():
     prophet.verify()
 
 
+Call = TypedDict("Call", {"name": str, "args": tuple[Any, ...], "kwargs": dict[str, Any]})
+Prediction = TypedDict("Prediction", {"type": Literal["SHOULD_BE_CALLED"], "number": Optional[int], "call": Call})
+
+
 class Prophecy:
-    def __init__(self, cls):
+    """
+    Rough outline how it works:
+    1. Test run pass:
+    - Remember calls (all calls should also have ``.should_be_called()`` after).
+    - Remember predictions.
+    - Associate return values with calls.
+
+    2. Core function pass:
+    - Remember calls.
+    - Use return values from the first pass.
+
+    3. Verification pass:
+    - Ensure all predicted calls actually happened.
+    """
+
+    subject: Type
+
+    def __init__(self, cls: Type):
         self.subject = cls
-        self.predictions: list[dict] = []
+        self.predictions: list[Prediction] = []
-        self.calls: list[dict] = []
+        self.calls: list[Call] = []
         self.return_values: dict[str, Any] = {}
-        self.should_call: Optional[dict] = None
+        self.should_call: Optional[Call] = None
 
     def __getattr__(self, attr: str):
         if not hasattr(self.subject, attr):
             raise AttributeError(f"Prophecy {self.subject} has no attribute {attr}")
 
-        def decorate(*args, **kwargs):
+        # It also returns `Any` but it only happens during `subject.xxx` call.
-            call = {"name": attr, "args": args, "kwargs": kwargs}
+        def decorate(*args: Any, **kwargs: Any) -> Self:
+            """Remember a call."""
+            call: Call = {"name": attr, "args": args, "kwargs": kwargs}
             # Ensure that signature is valid
             getattr(self.subject, attr)(*args, **kwargs)
             key = json.dumps(call, sort_keys=True)
@@ -265,18 +288,20 @@ class Prophecy:
 
         return decorate
 
-    def should_be_called(self, number=None):
+    def should_be_called(self, number: Optional[int] = None) -> Self:
+        """Predict the last added call."""
         self.should_call = self.calls.pop()
         self.predictions.append({"type": "SHOULD_BE_CALLED", "number": number, "call": self.should_call})
         return self
 
     def will_return(self, value: Any) -> Self:
+        """Remember a return value for the last predicted call."""
         key = json.dumps(self.should_call, sort_keys=True)
         self.return_values[key] = value
         return self
 
     def verify(self) -> None:
-        predicted_calls = []
+        predicted_calls: list[Call] = []
         for prediction in self.predictions:
             predicted_calls.append(prediction["call"])
             if prediction["type"] == "SHOULD_BE_CALLED":
@@ -285,7 +310,7 @@ class Prophecy:
             if call not in predicted_calls:
                 raise Exception(f"Unpredicted call: {call}")
 
-    def verify_should_be_called(self, prediction: dict) -> None:
+    def verify_should_be_called(self, prediction: Prediction) -> None:
         if prediction["number"]:
             count = self.calls.count(prediction["call"])
             if count != prediction["number"]: