Fix annotation data type coersion by Pydantic

Author: whistler

labelbox/data/annotation_types/data/audio.py

@@ -1,5 +1,7 @@
+from labelbox.typing_imports import Literal
+from labelbox.utils import _NoCoercionMixin
 from .base_data import BaseData
 
 
-class AudioData(BaseData):
-    ...
+class AudioData(BaseData, _NoCoercionMixin):
+    class_name: Literal["AudioData"] = "AudioData"

labelbox/data/annotation_types/data/conversation.py

@@ -1,5 +1,7 @@
+from labelbox.typing_imports import Literal
+from labelbox.utils import _NoCoercionMixin
 from .base_data import BaseData
 
 
 class ConversationData(BaseData):
-    ...
+    class_name: Literal["ConversationData"] = "ConversationData"

labelbox/data/annotation_types/data/dicom.py

@@ -1,5 +1,7 @@
+from labelbox.typing_imports import Literal
+from labelbox.utils import _NoCoercionMixin
 from .base_data import BaseData
 
 
-class DicomData(BaseData):
-    ...
+class DicomData(BaseData, _NoCoercionMixin):
+    class_name: Literal["DicomData"] = "DicomData"

labelbox/data/annotation_types/data/document.py

@@ -1,5 +1,7 @@
+from labelbox.typing_imports import Literal
+from labelbox.utils import _NoCoercionMixin
 from .base_data import BaseData
 
 
-class DocumentData(BaseData):
-    ...
+class DocumentData(BaseData, _NoCoercionMixin):
+    class_name: Literal["DocumentData"] = "DocumentData"

labelbox/data/annotation_types/data/html.py

@@ -1,5 +1,7 @@
+from labelbox.typing_imports import Literal
+from labelbox.utils import _NoCoercionMixin
 from .base_data import BaseData
 
 
-class HTMLData(BaseData):
-    ...
+class HTMLData(BaseData, _NoCoercionMixin):
+    class_name: Literal["HTMLData"] = "HTMLData"

labelbox/data/annotation_types/label.py

@@ -11,12 +11,15 @@
                          VideoClassificationAnnotation, VideoObjectAnnotation,
                          DICOMObjectAnnotation)
 from .classification import ClassificationAnswer
-from .data import DicomData, VideoData, TextData, ImageData
+from .data import AudioData, ConversationData, DicomData, DocumentData, HTMLData, ImageData, MaskData, TextData, VideoData
 from .geometry import Mask
 from .metrics import ScalarMetric, ConfusionMatrixMetric
 from .types import Cuid
 from ..ontology import get_feature_schema_lookup
 
+DataType = Union[AudioData, ConversationData, DicomData, DocumentData, HTMLData,
+                 ImageData, MaskData, TextData, VideoData]
+
 
 class Label(BaseModel):
     """Container for holding data and annotations
@@ -38,7 +41,7 @@ class Label(BaseModel):
         extra: additional context
     """
     uid: Optional[Cuid] = None
-    data: Union[VideoData, ImageData, TextData, TiledImageData]
+    data: DataType
     annotations: List[Union[ClassificationAnnotation, ObjectAnnotation,
                             ScalarMetric, ConfusionMatrixMetric]] = []
     extra: Dict[str, Any] = {}

labelbox/typing_imports.py

@@ -0,0 +1,10 @@
+"""
+This module imports types that differ across python versions, so other modules 
+don't have to worry about where they should be imported from.
+"""
+
+import sys
+if sys.version_info >= (3, 8):
+    from typing import Literal
+else:
+    from typing_extensions import Literal

labelbox/utils.py

@@ -35,3 +35,26 @@ class _CamelCaseMixin(BaseModel):
     class Config:
         allow_population_by_field_name = True
         alias_generator = camel_case
+
+
+class _NoCoercionMixin:
+    """
+    When using Unions in type annotations, pydantic will try to coerce the type
+    of the object to the type of the first Union member. Which results in
+    uninteded behavior.
+
+    This mixin uses a class_name discriminator field to prevent pydantic from
+    corecing the type of the object. Add a class_name field to the class you 
+    want to discrimniate and use this mixin class to remove the discriminator
+    when serializing the object.
+
+    Example:
+        class ConversationData(BaseData, _NoCoercionMixin):
+            class_name = Literal["ConversationData"] = "ConversationData"
+
+    """
+
+    def dict(self, *args, **kwargs):
+        res = super().dict(*args, **kwargs)
+        res.pop('class_name')
+        return res

tests/data/annotation_types/test_label.py

@@ -1,5 +1,6 @@
 import numpy as np
 
+import labelbox.types as lb_types
 from labelbox import OntologyBuilder, Tool, Classification as OClassification, Option
 from labelbox.data.annotation_types import (ClassificationAnswer, Radio, Text,
                                             ClassificationAnnotation,
@@ -181,3 +182,14 @@ def test_schema_assignment_confidence():
                   ])
 
     assert label.annotations[0].confidence == 0.914
+
+
+def test_initialize_label_no_coercion():
+    global_key = 'global-key'
+    ner_annotation = lb_types.ObjectAnnotation(
+        name="ner",
+        value=lb_types.ConversationEntity(start=0, end=8, message_id="4"))
+    label = Label(data=lb_types.ConversationData(global_key=global_key),
+                  annotations=[ner_annotation])
+    assert isinstance(label.data, lb_types.ConversationData)
+    assert label.data.global_key == global_key