From 34efa9f459d0c3d5ea52d4e1d9e1f3d2470f3fb9 Mon Sep 17 00:00:00 2001 From: Andrej730 Date: Fri, 25 Apr 2025 11:32:48 +0500 Subject: [PATCH] typing --- .../bonsai/bim/module/geometry/operator.py | 20 ++++++--- src/bonsai/bonsai/core/drawing.py | 2 +- src/bonsai/bonsai/tool/geometry.py | 10 +++-- src/bonsai/bonsai/tool/ifc.py | 4 +- src/bonsai/test/core/bootstrap.py | 45 ++++++++++++++----- 5 files changed, 59 insertions(+), 22 deletions(-) diff --git a/src/bonsai/bonsai/bim/module/geometry/operator.py b/src/bonsai/bonsai/bim/module/geometry/operator.py index 28f1ea0417..6597ced90f 100644 --- a/src/bonsai/bonsai/bim/module/geometry/operator.py +++ b/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( diff --git a/src/bonsai/bonsai/core/drawing.py b/src/bonsai/bonsai/core/drawing.py index 9687b3d876..a8c620dce3 100644 --- a/src/bonsai/bonsai/core/drawing.py +++ b/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(): diff --git a/src/bonsai/bonsai/tool/geometry.py b/src/bonsai/bonsai/tool/geometry.py index 645833baa0..2b9c87d7e6 100644 --- a/src/bonsai/bonsai/tool/geometry.py +++ b/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 - ) -> dict: + cls, + 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) diff --git a/src/bonsai/bonsai/tool/ifc.py b/src/bonsai/bonsai/tool/ifc.py index 2f228cb4f1..ba856323a3 100644 --- a/src/bonsai/bonsai/tool/ifc.py +++ b/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. diff --git a/src/bonsai/test/core/bootstrap.py b/src/bonsai/test/core/bootstrap.py index 9b89b81f6e..1f249e22b7 100644 --- a/src/bonsai/test/core/bootstrap.py +++ b/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.calls: list[dict] = [] + self.predictions: list[Prediction] = [] + 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): - call = {"name": attr, "args": args, "kwargs": kwargs} + # It also returns `Any` but it only happens during `subject.xxx` call. + 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"]: