Migrate mmc classes

Author: Val Brodsky

libs/labelbox/src/labelbox/data/annotation_types/mmc.py

@@ -1,7 +1,8 @@
 from abc import ABC
 from typing import ClassVar, List, Union
 
-from labelbox import pydantic_compat
+from pydantic import field_validator
+
 from labelbox.utils import _CamelCaseMixin
 from labelbox.data.annotation_types.annotation import BaseAnnotation
 
@@ -33,12 +34,15 @@ class MessageRankingTask(_BaseMessageEvaluationTask):
     format: ClassVar[str] = "message-ranking"
     ranked_messages: List[OrderedMessageInfo]
 
-    @pydantic_compat.validator("ranked_messages")
+    @field_validator("ranked_messages")
     def _validate_ranked_messages(cls, v: List[OrderedMessageInfo]):
         if not {msg.order for msg in v} == set(range(1, len(v) + 1)):
-            raise ValueError("Messages must be ordered by unique and consecutive natural numbers starting from 1")
+            raise ValueError(
+                "Messages must be ordered by unique and consecutive natural numbers starting from 1"
+            )
         return v
 
 
 class MessageEvaluationTaskAnnotation(BaseAnnotation):
-    value: Union[MessageSingleSelectionTask, MessageMultiSelectionTask, MessageRankingTask]
+    value: Union[MessageSingleSelectionTask, MessageMultiSelectionTask,
+                 MessageRankingTask]

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

@@ -2,17 +2,18 @@
 
 from labelbox.utils import _CamelCaseMixin
 
-from .base import DataRow, NDAnnotation
+from .base import _SubclassRegistryBase, DataRow, NDAnnotation
 from ...annotation_types.types import Cuid
 from ...annotation_types.mmc import MessageSingleSelectionTask, MessageMultiSelectionTask, MessageRankingTask, MessageEvaluationTaskAnnotation
 
 
 class MessageTaskData(_CamelCaseMixin):
     format: str
-    data: Union[MessageSingleSelectionTask, MessageMultiSelectionTask, MessageRankingTask]
+    data: Union[MessageSingleSelectionTask, MessageMultiSelectionTask,
+                MessageRankingTask]
 
 
-class NDMessageTask(NDAnnotation):
+class NDMessageTask(NDAnnotation, _SubclassRegistryBase):
 
     message_evaluation_task: MessageTaskData
 
@@ -26,17 +27,13 @@ def to_common(self) -> MessageEvaluationTaskAnnotation:
 
     @classmethod
     def from_common(
-        cls,
-        annotation: MessageEvaluationTaskAnnotation,
-        data: Any#Union[ImageData, TextData],
+            cls,
+            annotation: MessageEvaluationTaskAnnotation,
+            data: Any  #Union[ImageData, TextData],
     ) -> "NDMessageTask":
-        return cls(
-            uuid=str(annotation._uuid),
-            name=annotation.name,
-            schema_id=annotation.feature_schema_id,
-            data_row=DataRow(id=data.uid, global_key=data.global_key),
-            message_evaluation_task=MessageTaskData(
-                format=annotation.value.format,
-                data=annotation.value
-            )
-        )
+        return cls(uuid=str(annotation._uuid),
+                   name=annotation.name,
+                   schema_id=annotation.feature_schema_id,
+                   data_row=DataRow(id=data.uid, global_key=data.global_key),
+                   message_evaluation_task=MessageTaskData(
+                       format=annotation.value.format, data=annotation.value))

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

@@ -3,13 +3,12 @@
 import pytest
 
 from labelbox.data.serialization import NDJsonConverter
-from labelbox.pydantic_compat import ValidationError
 
 
 def test_message_task_annotation_serialization():
     with open('tests/data/assets/ndjson/mmc_import.json', 'r') as file:
         data = json.load(file)
-    
+
     deserialized = list(NDJsonConverter.deserialize(data))
     reserialized = list(NDJsonConverter.serialize(deserialized))
 
@@ -19,9 +18,12 @@ def test_message_task_annotation_serialization():
 def test_mesage_ranking_task_wrong_order_serialization():
     with open('tests/data/assets/ndjson/mmc_import.json', 'r') as file:
         data = json.load(file)
-    
-    some_ranking_task = next(task for task in data if task["messageEvaluationTask"]["format"] == "message-ranking")
-    some_ranking_task["messageEvaluationTask"]["data"]["rankedMessages"][0]["order"] = 3
 
-    with pytest.raises(ValidationError):
+    some_ranking_task = next(
+        task for task in data
+        if task["messageEvaluationTask"]["format"] == "message-ranking")
+    some_ranking_task["messageEvaluationTask"]["data"]["rankedMessages"][0][
+        "order"] = 3
+
+    with pytest.raises(ValueError):
         list(NDJsonConverter.deserialize([some_ranking_task]))