From a0792a4b9398abda7643243f2aa693a9287e9664 Mon Sep 17 00:00:00 2001
From: Andrej730 <azhilenkov@gmail.com>
Date: Wed, 21 Aug 2024 18:15:46 +0500
Subject: [PATCH] typing

---
 src/bonsai/bonsai/bim/ifc.py        |  5 ++++-
 src/bonsai/bonsai/bim/import_ifc.py |  4 +++-
 src/bonsai/bonsai/tool/collector.py |  1 +
 src/bonsai/bonsai/tool/geometry.py  | 16 ++++++++--------
 src/bonsai/bonsai/tool/ifc.py       |  6 +++++-
 src/bonsai/bonsai/tool/loader.py    |  6 +++++-
 6 files changed, 26 insertions(+), 12 deletions(-)

diff --git a/src/bonsai/bonsai/bim/ifc.py b/src/bonsai/bonsai/bim/ifc.py
index a6199c433b..f2f950e6ea 100644
--- a/src/bonsai/bonsai/bim/ifc.py
+++ b/src/bonsai/bonsai/bim/ifc.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 os
 import bpy
 import uuid
@@ -228,7 +229,9 @@ class IfcStore:
             IfcStore.commit_link_element(data)
 
     @staticmethod
-    def link_element(element: ifcopenshell.entity_instance, obj: IFC_CONNECTED_TYPE) -> None:
+    def link_element(
+        element: ifcopenshell.entity_instance, obj: Union[IFC_CONNECTED_TYPE, tool.Geometry.TYPES_WITH_MESH_PROPERTIES]
+    ) -> None:
         # Please use tool.Ifc.link() instead of this method. We want to
         # refactor this class and deprecate usage of IfcStore in favour of
         # tools.
diff --git a/src/bonsai/bonsai/bim/import_ifc.py b/src/bonsai/bonsai/bim/import_ifc.py
index 1d176b5bd3..6dd8d2e940 100644
--- a/src/bonsai/bonsai/bim/import_ifc.py
+++ b/src/bonsai/bonsai/bim/import_ifc.py
@@ -785,7 +785,9 @@ class IfcImporter:
             self.set_matrix_world(obj, tool.Loader.apply_blender_offset_to_matrix_world(obj, placement_matrix))
             self.link_element(product, obj)
 
-    def get_pointcloud_representation(self, product):
+    def get_pointcloud_representation(
+        self, product: ifcopenshell.entity_instance
+    ) -> Union[ifcopenshell.entity_instance, None]:
         if hasattr(product, "Representation") and hasattr(product.Representation, "Representations"):
             representations = product.Representation.Representations
         elif hasattr(product, "RepresentationMaps") and hasattr(product.RepresentationMaps, "RepresentationMaps"):
diff --git a/src/bonsai/bonsai/tool/collector.py b/src/bonsai/bonsai/tool/collector.py
index c5e1e8344f..db39675d75 100644
--- a/src/bonsai/bonsai/tool/collector.py
+++ b/src/bonsai/bonsai/tool/collector.py
@@ -40,6 +40,7 @@ class Collector(bonsai.core.tool.Collector):
                     users_collection.objects.unlink(obj)
 
         element = tool.Ifc.get_entity(obj)
+        assert element
 
         if element.is_a("IfcGridAxis"):
             element = (element.PartOfU or element.PartOfV or element.PartOfW)[0]
diff --git a/src/bonsai/bonsai/tool/geometry.py b/src/bonsai/bonsai/tool/geometry.py
index 04dc9bb923..910b921550 100644
--- a/src/bonsai/bonsai/tool/geometry.py
+++ b/src/bonsai/bonsai/tool/geometry.py
@@ -663,15 +663,15 @@ class Geometry(bonsai.core.tool.Geometry):
             return False
         return isinstance(data, supported_types)
 
+    TYPES_WITH_MESH_PROPERTIES = Union[
+        bpy.types.Mesh,
+        bpy.types.Curve,
+        bpy.types.Camera,
+        bpy.types.PointLight,
+    ]
+
     @classmethod
-    def has_mesh_properties(cls, data: Union[bpy.types.ID, None]) -> TypeIs[
-        Union[
-            bpy.types.Mesh,
-            bpy.types.Curve,
-            bpy.types.Camera,
-            bpy.types.PointLight,
-        ]
-    ]:
+    def has_mesh_properties(cls, data: Union[bpy.types.ID, None]) -> TypeIs[TYPES_WITH_MESH_PROPERTIES]:
         supported_types = (
             bpy.types.Mesh,
             bpy.types.Curve,
diff --git a/src/bonsai/bonsai/tool/ifc.py b/src/bonsai/bonsai/tool/ifc.py
index 7e3675a966..af7f3a3bf9 100644
--- a/src/bonsai/bonsai/tool/ifc.py
+++ b/src/bonsai/bonsai/tool/ifc.py
@@ -175,7 +175,11 @@ class Ifc(bonsai.core.tool.Ifc):
             bonsai.bim.handler.subscribe_to(obj, "name", bonsai.bim.handler.name_callback)
 
     @classmethod
-    def link(cls, element: ifcopenshell.entity_instance, obj: IFC_CONNECTED_TYPE) -> None:
+    def link(
+        cls,
+        element: ifcopenshell.entity_instance,
+        obj: Union[IFC_CONNECTED_TYPE, tool.Geometry.TYPES_WITH_MESH_PROPERTIES],
+    ) -> None:
         IfcStore.link_element(element, obj)
 
     @classmethod
diff --git a/src/bonsai/bonsai/tool/loader.py b/src/bonsai/bonsai/tool/loader.py
index 96fef5d8da..10fba187fb 100644
--- a/src/bonsai/bonsai/tool/loader.py
+++ b/src/bonsai/bonsai/tool/loader.py
@@ -98,7 +98,11 @@ class Loader(bonsai.core.tool.Loader):
         return "{}/{}".format(element.is_a(), getattr(element, "Name", "None"))
 
     @classmethod
-    def link_mesh(cls, shape, mesh: OBJECT_DATA_TYPE) -> None:
+    def link_mesh(
+        cls,
+        shape: Union[ifcopenshell.geom.ShapeElementType, ifcopenshell.geom.ShapeType],
+        mesh: tool.Geometry.TYPES_WITH_MESH_PROPERTIES,
+    ) -> None:
         geometry = shape.geometry if hasattr(shape, "geometry") else shape
         if "-" in geometry.id:
             mesh.BIMMeshProperties.ifc_definition_id = int(geometry.id.split("-")[0])