typing

src/ifcpatch/ifcpatch/recipes/ResetSpatialElementLocations.py

@@ -16,9 +16,12 @@
 # You should have received a copy of the GNU Lesser General Public License
 # along with IfcPatch.  If not, see <http://www.gnu.org/licenses/>.
 
+import ifcopenshell
+import logging
+
 
 class Patcher:
-    def __init__(self, file, logger, ifc_class="IfcSite"):
+    def __init__(self, file: ifcopenshell.file, logger: logging.Logger, ifc_class: str = "IfcSite"):
         """Resets the location of a spatial element to 0,0,0
 
         Another more specialised patch to fix incorrect coordinate usage is to
@@ -26,7 +29,6 @@ class Patcher:
         to 0,0,0.
 
         :param ifc_class: The class of spatial element to reset coordinates for.
-        :type ifc_class: str
 
         Example:
 
@@ -39,13 +41,15 @@ class Patcher:
         self.logger = logger
         self.ifc_class = ifc_class
 
-    def patch(self):
+    def patch(self) -> None:
         project = self.file.by_type("IfcProject")[0]
         spatial_elements = self.find_decomposed_ifc_class(project, self.ifc_class)
         for spatial_element in spatial_elements:
             self.patch_placement_to_origin(spatial_element)
 
-    def find_decomposed_ifc_class(self, element, ifc_class):
+    def find_decomposed_ifc_class(
+        self, element: ifcopenshell.entity_instance, ifc_class: str
+    ) -> list[ifcopenshell.entity_instance]:
         results = []
         rel_aggregates = element.IsDecomposedBy
         if not rel_aggregates:
@@ -57,7 +61,7 @@ class Patcher:
                 results.extend(self.find_decomposed_ifc_class(part, ifc_class))
         return results
 
-    def patch_placement_to_origin(self, element):
+    def patch_placement_to_origin(self, element: ifcopenshell.entity_instance) -> None:
         element.ObjectPlacement.RelativePlacement.Location.Coordinates = (0.0, 0.0, 0.0)
         if element.ObjectPlacement.RelativePlacement.Axis:
             element.ObjectPlacement.RelativePlacement.Axis.DirectionRatios = (0.0, 0.0, 1.0)