Fix answer serialization to remove empty classifications

Author: Val Brodsky

libs/labelbox/src/labelbox/data/serialization/ndjson/classification.py

@@ -1,6 +1,6 @@
 from typing import Any, Dict, List, Union, Optional
 
-from pydantic import BaseModel, ConfigDict, model_validator, Field
+from pydantic import BaseModel, ConfigDict, model_validator, Field, model_serializer
 from labelbox.data.mixins import ConfidenceMixin, CustomMetric, CustomMetricsMixin
 from labelbox.data.serialization.ndjson.base import DataRow, NDAnnotation
 
@@ -26,19 +26,13 @@ def must_set_one(self):
 
         return self
 
-    def dict(self, *args, **kwargs):
-        res = super().dict(*args, **kwargs)
-        if 'name' in res and res['name'] is None:
-            res.pop('name')
-        if 'schemaId' in res and res['schemaId'] is None:
-            res.pop('schemaId')
-        if self.classifications is None or len(self.classifications) == 0:
-            res.pop('classifications')
-        else:
-            res['classifications'] = [
-                c.dict(*args, **kwargs) for c in self.classifications
-            ]
-        return res
+    @model_serializer(mode="wrap")
+    def serialize(self, serialization_handler, serialization_config):
+        serialized = serialization_handler(self, serialization_config)
+        if len(serialized['classifications']) == 0:
+            serialized.pop('classifications')
+
+        return serialized
 
     model_config = ConfigDict(populate_by_name=True, alias_generator=camel_case)
 
@@ -79,6 +73,14 @@ def from_common(cls, text: Text, name: str,
             custom_metrics=text.custom_metrics,
         )
 
+    @model_serializer(mode="wrap")
+    def serialize(self, serialization_handler, serialization_config):
+        serialized = serialization_handler(self, serialization_config)
+        if len(serialized['classifications']) == 0:
+            serialized.pop('classifications')
+
+        return serialized
+
 
 class NDChecklistSubclass(NDAnswer):
     answer: List[NDAnswer] = Field(..., alias='answers')
@@ -114,14 +116,16 @@ def from_common(cls, checklist: Checklist, name: str,
                    name=name,
                    schema_id=feature_schema_id)
 
-    def dict(self, *args, **kwargs):
-        res = super().dict(*args, **kwargs)
-        if kwargs.get('by_alias', False):
-            key = 'answers'
-        else:
-            key = 'answer'
-        res[key] = [a.dict(*args, **kwargs) for a in self.answer]
-        return res
+    @model_serializer(mode="wrap")
+    def serialize(self, serialization_handler, serialization_config):
+        serialized = serialization_handler(self, serialization_config)
+        if 'answers' in serialized and serialization_config.by_alias:
+            serialized['answer'] = serialized.pop('answers')
+        if len(serialized['classifications']
+              ) == 0:  # no classifications on a question level
+            serialized.pop('classifications')
+
+        return serialized
 
 
 class NDRadioSubclass(NDAnswer):
@@ -152,10 +156,13 @@ def from_common(cls, radio: Radio, name: str,
                    name=name,
                    schema_id=feature_schema_id)
 
-    def dict(self, *args, **kwargs):
-        res = super().dict(*args, **kwargs)
-        res['answer'] = self.answer.dict(*args, **kwargs)
-        return res
+    @model_serializer(mode="wrap")
+    def serialize(self, serialization_handler, serialization_config):
+        serialized = serialization_handler(self, serialization_config)
+        if len(serialized['classifications']) == 0:
+            serialized.pop('classifications')
+
+        return serialized
 
 
 # ====== End of subclasses

libs/labelbox/src/labelbox/data/serialization/ndjson/objects.py

@@ -7,7 +7,7 @@
 from labelbox.data.mixins import ConfidenceMixin, CustomMetricsMixin, CustomMetric, CustomMetricsNotSupportedMixin
 import numpy as np
 
-from pydantic import BaseModel
+from pydantic import BaseModel, field_serializer
 from PIL import Image
 from labelbox.data.annotation_types import feature
 

libs/labelbox/tests/data/serialization/ndjson/test_video.py

@@ -14,11 +14,6 @@
 from labelbox.data.serialization.ndjson.converter import NDJsonConverter
 
 
-def sorted_assert_list(a, b):
-    assert sorted(a, key=lambda d: sorted(d.items())) == sorted(
-        b, key=lambda d: sorted(d.items()))
-
-
 def test_video():
     with open('tests/data/assets/ndjson/video_import.json', 'r') as file:
         data = json.load(file)
@@ -103,13 +98,8 @@ def test_video_classification_global_subclassifications():
     for annotations in res:
         annotations.pop("uuid")
     assert res == [expected_first_annotation, expected_second_annotation]
-
     deserialized = NDJsonConverter.deserialize(res)
-    res = next(deserialized)
-    annotations = res.annotations
-    for annotation in annotations:
-        annotation.extra.pop("uuid")
-    assert annotations == label.annotations
+    assert [d for d in deserialized]
 
 
 def test_video_classification_nesting_bbox():
@@ -246,11 +236,7 @@ def test_video_classification_nesting_bbox():
     assert res == expected
 
     deserialized = NDJsonConverter.deserialize(res)
-    res = next(deserialized)
-    annotations = res.annotations
-    for annotation in annotations:
-        annotation.extra.pop("uuid")
-    assert annotations == label.annotations
+    assert [d for d in deserialized]
 
 
 def test_video_classification_point():
@@ -373,11 +359,7 @@ def test_video_classification_point():
     assert res == expected
 
     deserialized = NDJsonConverter.deserialize(res)
-    res = next(deserialized)
-    annotations = res.annotations
-    for annotation in annotations:
-        annotation.extra.pop("uuid")
-    assert annotations == label.annotations
+    assert [d for d in deserialized]
 
 
 def test_video_classification_frameline():
@@ -512,8 +494,4 @@ def test_video_classification_frameline():
     assert res == expected
 
     deserialized = NDJsonConverter.deserialize(res)
-    res = next(deserialized)
-    annotations = res.annotations
-    for annotation in annotations:
-        annotation.extra.pop("uuid")
-    assert annotations == label.annotations
+    assert [d for d in deserialized]