typing

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

@@ -17,14 +17,15 @@
 # along with IfcOpenShell.  If not, see <http://www.gnu.org/licenses/>.
 
 import ifcopenshell.util.element
+from typing import Optional
 
 
-def get_references(element, should_inherit=True):
+def get_references(element: ifcopenshell.entity_instance, should_inherit=True) -> set[ifcopenshell.entity_instance]:
     results = set()
     if not element.is_a("IfcRoot"):
         if hasattr(element, "HasExternalReferences"):
             return {r.RelatingReference for r in element.HasExternalReferences or []}
-        elif hasattr(element, "HasExternalReference"): # Seriously, IFC?
+        elif hasattr(element, "HasExternalReference"):  # Seriously, IFC?
             return {r.RelatingReference for r in element.HasExternalReference or []}
     if should_inherit:
         element_type = ifcopenshell.util.element.get_type(element)
@@ -50,13 +51,13 @@ def get_references(element, should_inherit=True):
     return occurrence_results
 
 
-def get_classification(reference):
+def get_classification(reference: ifcopenshell.entity_instance) -> ifcopenshell.entity_instance:
     if reference.is_a("IfcClassification"):
         return reference
     return get_classification(reference.ReferencedSource)
 
 
-def get_inherited_references(reference):
+def get_inherited_references(reference: Optional[ifcopenshell.entity_instance]) -> list[ifcopenshell.entity_instance]:
     results = []
     while True:
         if not reference or reference.is_a("IfcClassification"):
@@ -64,4 +65,3 @@ def get_inherited_references(reference):
         results.append(reference)
         reference = reference.ReferencedSource
     return results
-