From c0aaebcbe96239d9db074e0596279886a1dc0edc Mon Sep 17 00:00:00 2001
From: Dion Moult <dion@thinkmoult.com>
Date: Mon, 27 Jan 2025 14:23:02 +1100
Subject: [PATCH] Fix bug where assigning units didn't necessarily assign to
 the project, and also didn't overwrite existing units of the same type.

---
 .../ifcopenshell/api/unit/assign_unit.py      | 22 +++++++---------
 .../test/api/unit/test_assign_unit.py         | 26 +++++++++++++++++--
 2 files changed, 34 insertions(+), 14 deletions(-)

diff --git a/src/ifcopenshell-python/ifcopenshell/api/unit/assign_unit.py b/src/ifcopenshell-python/ifcopenshell/api/unit/assign_unit.py
index 1b9c378c8d..4e2dc3f848 100644
--- a/src/ifcopenshell-python/ifcopenshell/api/unit/assign_unit.py
+++ b/src/ifcopenshell-python/ifcopenshell/api/unit/assign_unit.py
@@ -42,9 +42,7 @@ def assign_unit(
     :param units: A list of units to assign as project defaults. See
         ifcopenshell.api.unit.add_si_unit, unit.add_conversion_based_unit,
         and unit.add_monetary_unit for information on how to create units.
-    :type units: list[ifcopenshell.entity_instance],optional
     :return: The IfcUnitAssignment element
-    :rtype: ifcopenshell.entity_instance
 
     Example:
 
@@ -95,22 +93,22 @@ class Usecase:
         return unit_assignment
 
     def get_unit_assignment(self) -> ifcopenshell.entity_instance:
-        unit_assignment = self.file.by_type("IfcUnitAssignment")
-        if unit_assignment:
-            unit_assignment = unit_assignment[0]
-            # TODO: handle unit rewriting, which is complicated
-        else:
+        if not (unit_assignment := ifcopenshell.util.unit.get_unit_assignment(self.file)):
             unit_assignment = self.file.createIfcUnitAssignment()
-            if self.file.schema == "IFC2X3":
-                self.file.by_type("IfcProject")[0].UnitsInContext = unit_assignment
-            else:
-                self.file.by_type("IfcContext")[0].UnitsInContext = unit_assignment
+            self.file.by_type("IfcProject")[0].UnitsInContext = unit_assignment
         return unit_assignment
 
     def assign_units(
         self, unit_assignment: ifcopenshell.entity_instance, new_units: list[ifcopenshell.entity_instance]
     ) -> None:
-        units = set(unit_assignment.Units or [])
+        new_unit_types = [u.UnitType if not u.is_a("IfcMonetaryUnit") else u.is_a() for u in new_units]
+        units = set(
+            [
+                u
+                for u in (unit_assignment.Units or [])
+                if u.is_a() not in new_unit_types and getattr(u, "UnitType", None) not in new_unit_types
+            ]
+        )
         for unit in new_units:
             units.add(unit)
         unit_assignment.Units = list(units)
diff --git a/src/ifcopenshell-python/test/api/unit/test_assign_unit.py b/src/ifcopenshell-python/test/api/unit/test_assign_unit.py
index d2d58eaf2c..21a136b257 100644
--- a/src/ifcopenshell-python/test/api/unit/test_assign_unit.py
+++ b/src/ifcopenshell-python/test/api/unit/test_assign_unit.py
@@ -24,7 +24,7 @@ class TestAssignUnit(test.bootstrap.IFC4):
     def test_run(self):
         project = self.file.createIfcProject()
         unit1 = ifcopenshell.api.unit.add_monetary_unit(self.file, currency="USD")
-        unit2 = ifcopenshell.api.unit.add_monetary_unit(self.file, currency="JPY")
+        unit2 = ifcopenshell.api.unit.add_si_unit(self.file, unit_type="LENGTHUNIT", prefix="MILLI")
         assignment = ifcopenshell.api.unit.assign_unit(self.file, units=[unit1, unit2])
         assert project.UnitsInContext == assignment
         assert assignment.is_a("IfcUnitAssignment")
@@ -34,7 +34,7 @@ class TestAssignUnit(test.bootstrap.IFC4):
     def test_assign_units_to_an_existing_assignment(self):
         project = self.file.createIfcProject()
         unit1 = ifcopenshell.api.unit.add_monetary_unit(self.file, currency="USD")
-        unit2 = ifcopenshell.api.unit.add_monetary_unit(self.file, currency="JPY")
+        unit2 = ifcopenshell.api.unit.add_si_unit(self.file, unit_type="LENGTHUNIT", prefix="MILLI")
         assignment1 = ifcopenshell.api.unit.assign_unit(self.file, units=[unit1])
         assignment2 = ifcopenshell.api.unit.assign_unit(self.file, units=[unit2])
         assert project.UnitsInContext == assignment1
@@ -42,6 +42,28 @@ class TestAssignUnit(test.bootstrap.IFC4):
         assert unit1 in assignment1.Units
         assert unit2 in assignment1.Units
 
+    def test_overwriting_an_existing_unit_type(self):
+        project = self.file.createIfcProject()
+        unit1 = ifcopenshell.api.unit.add_si_unit(self.file, unit_type="LENGTHUNIT", prefix="MILLI")
+        unit2 = ifcopenshell.api.unit.add_si_unit(self.file, unit_type="LENGTHUNIT", prefix="CENTI")
+        assignment1 = ifcopenshell.api.unit.assign_unit(self.file, units=[unit1])
+        assignment2 = ifcopenshell.api.unit.assign_unit(self.file, units=[unit2])
+        assert project.UnitsInContext == assignment1
+        assert assignment1 == assignment2
+        assert unit1 not in assignment1.Units
+        assert unit2 in assignment1.Units
+
+    def test_overwriting_an_existing_monetary_unit(self):
+        project = self.file.createIfcProject()
+        unit1 = ifcopenshell.api.unit.add_monetary_unit(self.file, currency="USD")
+        unit2 = ifcopenshell.api.unit.add_monetary_unit(self.file, currency="AUD")
+        assignment1 = ifcopenshell.api.unit.assign_unit(self.file, units=[unit1])
+        assignment2 = ifcopenshell.api.unit.assign_unit(self.file, units=[unit2])
+        assert project.UnitsInContext == assignment1
+        assert assignment1 == assignment2
+        assert unit1 not in assignment1.Units
+        assert unit2 in assignment1.Units
+
 
 class TestAssignUnitIFC2X3(test.bootstrap.IFC2X3, TestAssignUnit):
     pass