typing

src/ifcopenshell-python/ifcopenshell/api/cost/add_cost_item_quantity.py

@@ -17,10 +17,13 @@
 # along with IfcOpenShell.  If not, see <http://www.gnu.org/licenses/>.
 
 import ifcopenshell.api
+import ifcopenshell.util.unit
 
 
 def add_cost_item_quantity(
-    file: ifcopenshell.file, cost_item: ifcopenshell.entity_instance, ifc_class: str = "IfcQuantityCount"
+    file: ifcopenshell.file,
+    cost_item: ifcopenshell.entity_instance,
+    ifc_class: ifcopenshell.util.unit.QUANTITY_CLASS = "IfcQuantityCount",
 ) -> ifcopenshell.entity_instance:
     """Adds a new quantity associated with a cost item
 

src/ifcopenshell-python/ifcopenshell/api/cost/assign_cost_item_quantity.py

@@ -17,14 +17,14 @@
 # along with IfcOpenShell.  If not, see <http://www.gnu.org/licenses/>.
 
 import ifcopenshell.api
-from typing import Optional
+from typing import Any
 
 
 def assign_cost_item_quantity(
     file: ifcopenshell.file,
     cost_item: ifcopenshell.entity_instance,
     products: list[ifcopenshell.entity_instance],
-    prop_name: Optional[str] = "",
+    prop_name: str = "",
 ) -> None:
     """Adds a cost item quantity that is parametrically connected to a product
 
@@ -96,6 +96,9 @@ def assign_cost_item_quantity(
 
 
 class Usecase:
+    file: ifcopenshell.file
+    settings: dict[str, Any]
+
     def execute(self):
         if self.settings["prop_name"]:
             self.quantities = set(self.settings["cost_item"].CostQuantities or [])
@@ -113,7 +116,9 @@ class Usecase:
         else:
             self.update_cost_item_count()
 
-    def assign_cost_control(self, related_object, cost_item):
+    def assign_cost_control(
+        self, related_object: ifcopenshell.entity_instance, cost_item: ifcopenshell.entity_instance
+    ) -> ifcopenshell.entity_instance:
         return ifcopenshell.api.run(
             "control.assign_control",
             self.file,
@@ -121,12 +126,12 @@ class Usecase:
             relating_control=cost_item,
         )
 
-    def add_quantity_from_related_object(self, element):
+    def add_quantity_from_related_object(self, element: ifcopenshell.entity_instance) -> None:
         for relationship in element.IsDefinedBy:
             if relationship.is_a("IfcRelDefinesByProperties"):
                 self.add_quantity_from_qto(relationship.RelatingPropertyDefinition)
 
-    def add_quantity_from_qto(self, qto):
+    def add_quantity_from_qto(self, qto: ifcopenshell.entity_instance) -> None:
         if not qto.is_a("IfcElementQuantity"):
             return
         for prop in qto.Quantities:

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

@@ -337,6 +337,25 @@ unit_symbols = {
     "fahrenheit": "°F",
 }
 
+QUANTITY_CLASS = Literal[
+    "IfcQuantityCount",
+    "IfcQuantityNumber",
+    "IfcQuantityLength",
+    "IfcQuantityArea",
+    "IfcQuantityVolume",
+    "IfcQuantityWeight",
+    "IfcQuantityTime",
+    "IfcQuantityCount",
+]
+
+MEASURE_CLASS = Literal[
+    "IfcNumericMeasure",
+    "IfcLengthMeasure",
+    "IfcAreaMeasure",
+    "IfcVolumeMeasure",
+    "IfcMassMeasure",
+]
+
 
 def get_prefix(text):
     if text:
@@ -468,14 +487,14 @@ def get_property_unit(
         return units[0]
 
 
-def get_unit_measure_class(unit_type: str) -> str:
+def get_unit_measure_class(unit_type: str) -> MEASURE_CLASS:
     if unit_type == "USERDEFINED":
         # See https://github.com/buildingSMART/IFC4.3.x-development/issues/71
         return "IfcNumericMeasure"
     return "Ifc" + unit_type[0:-4].lower().capitalize() + "Measure"
 
 
-def get_measure_unit_type(measure_class: str) -> str:
+def get_measure_unit_type(measure_class: MEASURE_CLASS) -> str:
     if measure_class == "IfcNumericMeasure":
         # See https://github.com/buildingSMART/IFC4.3.x-development/issues/71
         return "USERDEFINED"
@@ -484,7 +503,7 @@ def get_measure_unit_type(measure_class: str) -> str:
     return measure_class.upper() + "UNIT"
 
 
-def get_symbol_measure_class(symbol: Optional[str] = None) -> str:
+def get_symbol_measure_class(symbol: Optional[str] = None) -> MEASURE_CLASS:
     # Dumb, but everybody gets it, unlike regex golf
     if not symbol:
         return "IfcNumericMeasure"
@@ -502,7 +521,7 @@ def get_symbol_measure_class(symbol: Optional[str] = None) -> str:
     return "IfcNumericMeasure"
 
 
-def get_symbol_quantity_class(symbol: Optional[str] = None) -> str:
+def get_symbol_quantity_class(symbol: Optional[str] = None) -> QUANTITY_CLASS:
     # Dumb, but everybody gets it, unlike regex golf
     if not symbol:
         return "IfcQuantityCount"