diff --git a/src/blenderbim/blenderbim/bim/module/classification/operator.py b/src/blenderbim/blenderbim/bim/module/classification/operator.py
index 84f2f89599..08020cef76 100644
--- a/src/blenderbim/blenderbim/bim/module/classification/operator.py
+++ b/src/blenderbim/blenderbim/bim/module/classification/operator.py
@@ -20,6 +20,7 @@ import bpy
 import json
 import ifcopenshell
 import ifcopenshell.api
+import ifcopenshell.util.element
 import blenderbim.tool as tool
 import blenderbim.bim.helper
 from blenderbim.bim.ifc import IfcStore
diff --git a/src/ifcopenshell-python/ifcopenshell/api/classification/add_classification.py b/src/ifcopenshell-python/ifcopenshell/api/classification/add_classification.py
index 0682ab7565..d59a4e9343 100644
--- a/src/ifcopenshell-python/ifcopenshell/api/classification/add_classification.py
+++ b/src/ifcopenshell-python/ifcopenshell/api/classification/add_classification.py
@@ -19,10 +19,11 @@
 import ifcopenshell
 import ifcopenshell.util.schema
 import ifcopenshell.util.date
+from typing import Union
 
 
 class Usecase:
-    def __init__(self, file, classification=None):
+    def __init__(self, file: ifcopenshell.file, classification: Union[str, ifcopenshell.entity_instance]):
         """Adds a new classification system to the project
 
         External classification systems such as Uniclass or Omniclass are
@@ -81,7 +82,7 @@ class Usecase:
             "classification": classification,
         }
 
-    def execute(self):
+    def execute(self) -> ifcopenshell.entity_instance:
         if isinstance(self.settings["classification"], str):
             classification = self.file.createIfcClassification(Name=self.settings["classification"])
             self.relate_to_project(classification)