From 7b62b035c9cd6099b256740d69347446e793dc20 Mon Sep 17 00:00:00 2001
From: Dion Moult <dion@thinkmoult.com>
Date: Mon, 20 Jan 2025 23:50:03 +1100
Subject: [PATCH] Fix #5964. Only get the unit assignment from the project to
 prevent ambiguity.

I really don't know why I didn't do this before.
---
 src/bonsai/bonsai/bim/import_ifc.py               | 4 ++--
 src/bonsai/bonsai/tool/unit.py                    | 9 ++++-----
 src/ifcopenshell-python/ifcopenshell/util/unit.py | 7 ++-----
 3 files changed, 8 insertions(+), 12 deletions(-)

diff --git a/src/bonsai/bonsai/bim/import_ifc.py b/src/bonsai/bonsai/bim/import_ifc.py
index 144fd4e5c1..87cebbb7f0 100644
--- a/src/bonsai/bonsai/bim/import_ifc.py
+++ b/src/bonsai/bonsai/bim/import_ifc.py
@@ -893,9 +893,9 @@ class IfcImporter:
         tool.Loader.set_unit_scale(self.unit_scale)
 
     def set_units(self):
-        if not (units := self.file.by_type("IfcUnitAssignment")):
+        if not (assignment := self.file.by_type("IfcProject")[0].UnitsInContext):
             return  # Geometry is optional in IFC
-        for unit in units[0].Units:
+        for unit in assignment.Units:
             if unit.is_a("IfcNamedUnit") and unit.UnitType == "LENGTHUNIT":
                 if unit.is_a("IfcSIUnit"):
                     bpy.context.scene.unit_settings.system = "METRIC"
diff --git a/src/bonsai/bonsai/tool/unit.py b/src/bonsai/bonsai/tool/unit.py
index 1ab1375268..37597cf95c 100644
--- a/src/bonsai/bonsai/tool/unit.py
+++ b/src/bonsai/bonsai/tool/unit.py
@@ -116,9 +116,9 @@ class Unit(bonsai.core.tool.Unit):
         for unit_class in ["IfcDerivedUnit", "IfcMonetaryUnit", "IfcNamedUnit"]:
             units += tool.Ifc.get().by_type(unit_class)
 
-        assigned_units = tool.Ifc.get().by_type("IfcUnitAssignment")
-        if assigned_units:
-            assigned_units = assigned_units[0].Units
+        assigned_units = []
+        if assignment := tool.Ifc.get().by_type("IfcProject")[0].UnitsInContext:
+            assigned_units = assignment.Units
 
         for unit in units:
             name = ""
@@ -162,8 +162,7 @@ class Unit(bonsai.core.tool.Unit):
 
     @classmethod
     def get_project_currency_unit(cls) -> Union[ifcopenshell.entity_instance, None]:
-        unit_assignments = tool.Ifc.get().by_type("IfcUnitAssignment")
-        for assignment in unit_assignments:
+        if assignment := tool.Ifc.get().by_type("IfcProject")[0].UnitsInContext:
             for unit in assignment.Units:
                 if unit.is_a("IfcMonetaryUnit"):
                     return unit
diff --git a/src/ifcopenshell-python/ifcopenshell/util/unit.py b/src/ifcopenshell-python/ifcopenshell/util/unit.py
index a19594f1b5..a8dd0dfdb4 100644
--- a/src/ifcopenshell-python/ifcopenshell/util/unit.py
+++ b/src/ifcopenshell-python/ifcopenshell/util/unit.py
@@ -404,9 +404,7 @@ def get_named_dimensions(name):
 
 
 def get_unit_assignment(ifc_file: ifcopenshell.file) -> Union[ifcopenshell.entity_instance, None]:
-    unit_assignments = ifc_file.by_type("IfcUnitAssignment")
-    if unit_assignments:
-        return unit_assignments[0]
+    return ifc_file.by_type("IfcProject")[0].UnitsInContext
 
 
 def get_project_unit(ifc_file: ifcopenshell.file, unit_type: str) -> Union[ifcopenshell.entity_instance, None]:
@@ -641,9 +639,8 @@ def calculate_unit_scale(ifc_file: ifcopenshell.file, unit_type: str = "LENGTHUN
     :param unit_type: The type of SI unit, defaults to "LENGTHUNIT"
     :returns: The scale factor
     """
-    if not ifc_file.by_type("IfcUnitAssignment"):
+    if not (units := ifc_file.by_type("IfcProject")[0].UnitsInContext):
         return 1
-    units = ifc_file.by_type("IfcUnitAssignment")[0]
     unit_scale = 1
     for unit in units.Units:
         if not hasattr(unit, "UnitType") or unit.UnitType != unit_type: