From 02ae838c4520a7069fbb50cc9b095bbc30afaab9 Mon Sep 17 00:00:00 2001
From: Andrej730 <azhilenkov@gmail.com>
Date: Tue, 16 Apr 2024 11:32:32 +0500
Subject: [PATCH] util.constraint - typing and docs

---
 .../ifcopenshell/util/constraint.py           | 47 ++++++++++++++-----
 1 file changed, 36 insertions(+), 11 deletions(-)

diff --git a/src/ifcopenshell-python/ifcopenshell/util/constraint.py b/src/ifcopenshell-python/ifcopenshell/util/constraint.py
index b2a38d6060..5df7f52846 100644
--- a/src/ifcopenshell-python/ifcopenshell/util/constraint.py
+++ b/src/ifcopenshell-python/ifcopenshell/util/constraint.py
@@ -18,20 +18,43 @@
 #
 #
 
-def get_constraints(product):
+import ifcopenshell
+from typing import Union
+
+
+def get_constraints(product: ifcopenshell.entity_instance) -> list[ifcopenshell.entity_instance]:
+    """
+    Retrieves the constraints assigned to the `product`.
+
+    :param product: The IFC element.
+    :type product: ifcopenshell.entity_instance.entity_instance
+    :return: List of assigned constraints.
+    :rtype: list[ifcopenshell.entity_instance.entity_instance]
+    """
     constraints = []
     for rel in product.HasAssociations or []:
         if rel.is_a("IfcRelAssociatesConstraint"):
             constraints.append(rel.RelatingConstraint)
     return constraints
 
-def get_metrics(constraint):
+
+def get_metrics(constraint: ifcopenshell.entity_instance) -> list[ifcopenshell.entity_instance]:
+    """
+    Retrieves the list of nested constraints for a IfcObjective `constraint`.
+
+    :param product: IfcObjective constraint.
+    :type product: ifcopenshell.entity_instance.entity_instance
+    :return: List of nested constraints.
+    :rtype: list[ifcopenshell.entity_instance.entity_instance]
+    """
+
     metrics = []
     for metric in constraint.BenchmarkValues or []:
         metrics.append(metric)
     return metrics
 
-def get_metric_reference(metric, is_deep=True):
+
+def get_metric_reference(metric: ifcopenshell.entity_instance, is_deep=True):
     def get_reference_Attribute(ref, path):
         if ref:
             if is_deep:
@@ -47,7 +70,10 @@ def get_metric_reference(metric, is_deep=True):
     reference = metric.ReferencePath
     return get_reference_Attribute(reference, "")
 
-def get_metric_constraints(resource, attribute):
+
+def get_metric_constraints(
+    resource: ifcopenshell.entity_instance, attribute
+) -> Union[list[ifcopenshell.entity_instance], None]:
     metrics = []
     for constraint in get_constraints(resource) or []:
         for metric in get_metrics(constraint) or []:
@@ -60,15 +86,14 @@ def get_metric_constraints(resource, attribute):
         return metrics
     return None
 
-def is_hard_constraint(metric):
-    if metric.ConstraintGrade == "HARD" and metric.Benchmark == "EQUALTO":
-        return True
 
-def is_attribute_locked(product, attribute):
+def is_hard_constraint(metric: ifcopenshell.entity_instance) -> bool:
+    return metric.ConstraintGrade == "HARD" and metric.Benchmark == "EQUALTO"
+
+
+def is_attribute_locked(product: ifcopenshell.entity_instance, attribute) -> bool:
     is_locked = False
-    metrics = get_metric_constraints(
-        product, attribute
-    )
+    metrics = get_metric_constraints(product, attribute)
     for metric in metrics or []:
         if is_hard_constraint(metric):
             is_locked = True