typing

src/ifcopenshell-python/ifcopenshell/util/representation.py

@@ -18,9 +18,16 @@
 
 import numpy as np
 import ifcopenshell
+import ifcopenshell.util.placement
+from typing import Optional, Union, TypedDict
 
 
-def get_context(ifc_file, context, subcontext=None, target_view=None):
+def get_context(
+    ifc_file: ifcopenshell.file,
+    context: str,
+    subcontext: Optional[str] = None,
+    target_view: Optional[str] = None,
+) -> Union[ifcopenshell.entity_instance, None]:
     if subcontext or target_view:
         elements = ifc_file.by_type("IfcGeometricRepresentationSubContext")
     else:
@@ -35,7 +42,12 @@ def get_context(ifc_file, context, subcontext=None, target_view=None):
         return element
 
 
-def is_representation_of_context(representation, context, subcontext=None, target_view=None):
+def is_representation_of_context(
+    representation: ifcopenshell.entity_instance,
+    context: Union[ifcopenshell.entity_instance, str],
+    subcontext: Optional[str] = None,
+    target_view: Optional[str] = None,
+) -> bool:
     if isinstance(context, ifcopenshell.entity_instance):
         return representation.ContextOfItems == context
 
@@ -52,11 +64,16 @@ def is_representation_of_context(representation, context, subcontext=None, targe
             and representation.ContextOfItems.ContextIdentifier == subcontext
             and representation.ContextOfItems.ContextType == context
         )
-    elif representation.ContextOfItems.ContextType == context:
-        return True
+
+    return representation.ContextOfItems.ContextType == context
 
 
-def get_representation(element, context, subcontext=None, target_view=None):
+def get_representation(
+    element: ifcopenshell.entity_instance,
+    context: Union[ifcopenshell.entity_instance, str],
+    subcontext: Optional[str] = None,
+    target_view: Optional[str] = None,
+) -> Union[ifcopenshell.entity_instance, None]:
     if element.is_a("IfcProduct") and element.Representation:
         for r in element.Representation.Representations:
             if is_representation_of_context(r, context, subcontext, target_view):
@@ -67,16 +84,30 @@ def get_representation(element, context, subcontext=None, target_view=None):
                 return r.MappedRepresentation
 
 
-def resolve_representation(representation):
+def resolve_representation(representation: ifcopenshell.entity_instance) -> ifcopenshell.entity_instance:
+    """Resolve possibly mapped representation.
+
+    :param representation: IfcRepresentation
+    :type representation: ifcopenshell.entity_instance.entity_instance
+    :return: Representation resolved from mappings
+    :rtype: ifcopenshell.entity_instance.entity_instance
+    """
     if len(representation.Items) == 1 and representation.Items[0].is_a("IfcMappedItem"):
         return resolve_representation(representation.Items[0].MappingSource.MappedRepresentation)
     return representation
 
 
-def resolve_items(representation, matrix=None):
+class ResolvedItemDict(TypedDict):
+    matrix: np.array
+    item: ifcopenshell.entity_instance
+
+
+def resolve_items(
+    representation: ifcopenshell.entity_instance, matrix: Optional[np.array] = None
+) -> list[ResolvedItemDict]:
     if matrix is None:
         matrix = np.eye(4)
-    results = []
+    results: list[ResolvedItemDict] = []
     for item in representation.Items or []:  # Be forgiving of invalid IFCs because Revit :(
         if item.is_a("IfcMappedItem"):
             rep_matrix = ifcopenshell.util.placement.get_mappeditem_transformation(item)
@@ -84,5 +115,5 @@ def resolve_items(representation, matrix=None):
                 rep_matrix = rep_matrix @ matrix.copy()
             results.extend(resolve_items(item.MappingSource.MappedRepresentation, rep_matrix))
         else:
-            results.append({"matrix": matrix.copy(), "item": item})
+            results.append(ResolvedItemDict(matrix=matrix.copy(), item=item))
     return results