Make Category immutable and add tests (#1206)

Author: gboeer

sahi/annotation.py

@@ -145,25 +145,20 @@ def __repr__(self):
         return f"BoundingBox: <{(self.minx, self.miny, self.maxx, self.maxy)}, w: {self.maxx - self.minx}, h: {self.maxy - self.miny}>"
 
 
+@dataclass(frozen=True)
 class Category:
     """
     Category of the annotation.
     """
 
-    def __init__(self, id=None, name=None):
-        """
-        Args:
-            id: int
-                ID of the object category
-            name: str
-                Name of the object category
-        """
-        if not isinstance(id, int):
+    id: int
+    name: str
+
+    def __post_init__(self):
+        if not isinstance(self.id, int):
             raise TypeError("id should be integer")
-        if not isinstance(name, str):
+        if not isinstance(self.name, str):
             raise TypeError("name should be string")
-        self.id = id
-        self.name = name
 
     def __repr__(self):
         return f"Category: <id: {self.id}, name: {self.name}>"

tests/test_annotation.py

@@ -63,6 +63,33 @@ def test_category(self):
         assert category.id == category_id
         assert category.name == category_name
 
+        # id must be int
+        with pytest.raises(TypeError):
+            Category(id="not-an-int", name="car")
+
+        # name must be str
+        with pytest.raises(TypeError):
+            Category(id=1, name=123)
+
+    def test_category_immutability(self):
+        import dataclasses
+
+        from sahi.annotation import Category
+
+        category = Category(id=5, name="person")
+
+        # Attempt to mutate the id directly
+        with pytest.raises(dataclasses.FrozenInstanceError):
+            category.id = 10
+
+        # Attempt to mutate the name directly
+        with pytest.raises(dataclasses.FrozenInstanceError):
+            category.name = "cat"
+
+        # Confirm the values remain unchanged
+        assert category.id == 5
+        assert category.name == "person"
+
     def test_mask(self):
         from sahi.annotation import Mask