typing

src/ifcopenshell-python/ifcopenshell/api/document/add_reference.py

@@ -16,9 +16,11 @@
 # You should have received a copy of the GNU Lesser General Public License
 # along with IfcOpenShell.  If not, see <http://www.gnu.org/licenses/>.
 
+import ifcopenshell
+
 
 class Usecase:
-    def __init__(self, file, information=None):
+    def __init__(self, file: ifcopenshell.file, information: ifcopenshell.entity_instance):
         """Creates a new reference to a document to assign to products
 
         A document may be associated with physical products, tasks, cost items,
@@ -65,7 +67,7 @@ class Usecase:
         self.file = file
         self.settings = {"information": information}
 
-    def execute(self):
+    def execute(self) -> ifcopenshell.entity_instance:
         if self.file.schema == "IFC2X3":
             reference = self.file.create_entity("IfcDocumentReference")
             if self.settings["information"]:

src/ifcopenshell-python/ifcopenshell/api/document/remove_reference.py

@@ -21,7 +21,7 @@ import ifcopenshell.util.element
 
 
 class Usecase:
-    def __init__(self, file, reference=None):
+    def __init__(self, file: ifcopenshell.file, reference: ifcopenshell.entity_instance):
         """Remove a document reference
 
         All associations with objects are removed.
@@ -42,7 +42,7 @@ class Usecase:
         self.file = file
         self.settings = {"reference": reference}
 
-    def execute(self):
+    def execute(self) -> None:
         for rel in self.settings["reference"].DocumentRefForObjects or []:
             history = rel.OwnerHistory
             self.file.remove(rel)