typing

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

@@ -18,9 +18,10 @@
 
 import numpy as np
 import ifcopenshell
+from typing import Literal, Iterable
 
 
-def a2p(o, z, x):
+def a2p(o: Iterable[float], z: Iterable[float], x: Iterable[float]) -> np.array:
     """Converts a location, X, and Z axis vector to a 4x4 transformation matrix
 
     IFC uses a right-handed coordinate system, so it is not necessary to
@@ -73,7 +74,7 @@ def get_axis2placement(placement: ifcopenshell.entity_instance) -> np.array:
     return a2p(o, z, x)
 
 
-def get_local_placement(placement):
+def get_local_placement(placement: ifcopenshell.entity_instance) -> np.array:
     """Parse a local placement into a 4x4 transformation matrix
 
     This is typically used to find the location and rotation of an element. The
@@ -107,7 +108,7 @@ def get_local_placement(placement):
     return np.dot(parent, get_axis2placement(placement.RelativePlacement))
 
 
-def get_cartesiantransformationoperator3d(inst):
+def get_cartesiantransformationoperator3d(inst: ifcopenshell.entity_instance) -> np.array:
     """Parses an IfcCartesianTransformationOperator into a 4x4 transformation matrix
 
     Note that in general you will not need to call this directly. See
@@ -152,7 +153,7 @@ def get_cartesiantransformationoperator3d(inst):
     return m4
 
 
-def get_mappeditem_transformation(item):
+def get_mappeditem_transformation(item: ifcopenshell.entity_instance) -> np.array:
     """Parse an IfcMappedItem into a 4x4 transformation matrix
 
     Mapped items take a representation with an origin and transform them with a
@@ -170,7 +171,7 @@ def get_mappeditem_transformation(item):
         return get_cartesiantransformationoperator3d(item.MappingTarget) @ m4
 
 
-def get_storey_elevation(storey):
+def get_storey_elevation(storey: ifcopenshell.entity_instance) -> float:
     """Get the Z elevation in project units of a buildling storey
 
     Building storeys store elevation in two possible locations: the Z value of
@@ -187,7 +188,7 @@ def get_storey_elevation(storey):
     return getattr(storey, "Elevation", 0.0) or 0.0
 
 
-def rotation(angle, axis, is_degrees=True):
+def rotation(angle: float, axis: Literal["X", "Y", "Z"], is_degrees=True) -> np.array:
     """Create a 4x4 numpy matrix representing an euler rotation
 
     :param angle: The angle of rotation