Refactored FeatureSchemaAttribute to be a dataclass

Author: Tim-Kerr

libs/labelbox/src/labelbox/schema/ontology.py

@@ -25,7 +25,10 @@
 from labelbox.schema.tool_building.tool_type_mapping import (
     map_tool_type_to_tool_cls,
 )
-from labelbox.schema.tool_building.types import FeatureSchemaAttributes
+from labelbox.schema.tool_building.types import (
+    FeatureSchemaAttribute,
+    FeatureSchemaAttributes,
+)
 import warnings
 
 
@@ -119,7 +122,15 @@ def from_dict(cls, dictionary: Dict[str, Any]) -> Dict[str, Any]:
                 for c in dictionary["classifications"]
             ],
             color=dictionary["color"],
-            attributes=dictionary.get("attributes", None),
+            attributes=[
+                FeatureSchemaAttribute(
+                    attributeName=attr["attributeName"],
+                    attributeValue=attr["attributeValue"],
+                )
+                for attr in dictionary.get("attributes", []) or []
+            ]
+            if dictionary.get("attributes")
+            else None,
         )
 
     def asdict(self) -> Dict[str, Any]:
@@ -133,7 +144,9 @@ def asdict(self) -> Dict[str, Any]:
             ],
             "schemaNodeId": self.schema_id,
             "featureSchemaId": self.feature_schema_id,
-            "attributes": self.attributes,
+            "attributes": [a.asdict() for a in self.attributes]
+            if self.attributes is not None
+            else None,
         }
 
     def add_classification(self, classification: Classification) -> None:

libs/labelbox/src/labelbox/schema/tool_building/classification.py

@@ -8,6 +8,7 @@
 from labelbox.schema.tool_building.types import (
     FeatureSchemaId,
     FeatureSchemaAttributes,
+    FeatureSchemaAttribute,
 )
 
 
@@ -112,7 +113,15 @@ def from_dict(cls, dictionary: Dict[str, Any]) -> "Classification":
             schema_id=dictionary.get("schemaNodeId", None),
             feature_schema_id=dictionary.get("featureSchemaId", None),
             scope=cls.Scope(dictionary.get("scope", cls.Scope.GLOBAL)),
-            attributes=dictionary.get("attributes", None),
+            attributes=[
+                FeatureSchemaAttribute(
+                    attributeName=attr["attributeName"],
+                    attributeValue=attr["attributeValue"],
+                )
+                for attr in dictionary.get("attributes", []) or []
+            ]
+            if dictionary.get("attributes")
+            else None,
         )
 
     def asdict(self, is_subclass: bool = False) -> Dict[str, Any]:
@@ -128,7 +137,7 @@ def asdict(self, is_subclass: bool = False) -> Dict[str, Any]:
             "options": [o.asdict() for o in self.options],
             "schemaNodeId": self.schema_id,
             "featureSchemaId": self.feature_schema_id,
-            "attributes": self.attributes
+            "attributes": [a.asdict() for a in self.attributes]
             if self.attributes is not None
             else None,
         }

libs/labelbox/src/labelbox/schema/tool_building/types.py

@@ -3,12 +3,22 @@
 from typing import TypedDict
 
 
-class FeatureSchemaAttribute(TypedDict):
+from dataclasses import dataclass
+
+
+@dataclass
+class FeatureSchemaAttribute:
     attributeName: str
     attributeValue: str
 
+    def asdict(self):
+        return {
+            "attributeName": self.attributeName,
+            "attributeValue": self.attributeValue,
+        }
+
 
-FeatureSchemaAttriubte = Annotated[FeatureSchemaAttribute, Field()]
+FeatureSchemaAttribute = Annotated[FeatureSchemaAttribute, Field()]
 
 FeatureSchemaId = Annotated[str, Field(min_length=25, max_length=25)]
 SchemaId = Annotated[str, Field(min_length=25, max_length=25)]

libs/labelbox/tests/unit/test_unit_ontology.py

@@ -36,9 +36,10 @@
             "color": "#FF0000",
             "tool": "rectangle",
             "attributes": [
-                FeatureSchemaAttribute(
-                    attriubteName="auto-ocr", attributeValue="true"
-                )
+                {
+                    "attributeName": "auto-ocr",
+                    "attributeValue": "true",
+                }
             ],
             "classifications": [
                 {
@@ -77,10 +78,10 @@
                         },
                     ],
                     "attributes": [
-                        FeatureSchemaAttribute(
-                            attributeName="requires-connection",
-                            attributeValue="true",
-                        )
+                        {
+                            "attributeName": "requires-connection",
+                            "attributeValue": "true",
+                        }
                     ],
                 },
                 {