[PLT-600] Converted SDK to use pydantic V2 (#1738)

Co-authored-by: Val Brodsky <[email protected]>

Author: Gabefire

libs/labelbox/pyproject.toml

@@ -5,7 +5,7 @@ description = "Labelbox Python API"
 authors = [{ name = "Labelbox", email = "[email protected]" }]
 dependencies = [
     "google-api-core>=1.22.1",
-    "pydantic>=1.8",
+    "pydantic>=2.0",
     "python-dateutil>=2.8.2, <2.10.0",
     "requests>=2.22.0",
     "strenum>=0.4.15",

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

@@ -29,7 +29,6 @@
 
 from .classification import Checklist
 from .classification import ClassificationAnswer
-from .classification import Dropdown
 from .classification import Radio
 from .classification import Text
 

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

@@ -7,6 +7,7 @@
 
 from labelbox.data.annotation_types.classification.classification import ClassificationAnnotation
 from .ner import DocumentEntity, TextEntity, ConversationEntity
+from typing import Optional
 
 
 class ObjectAnnotation(BaseAnnotation, ConfidenceMixin, CustomMetricsMixin):
@@ -29,4 +30,4 @@ class ObjectAnnotation(BaseAnnotation, ConfidenceMixin, CustomMetricsMixin):
     """
 
     value: Union[TextEntity, ConversationEntity, DocumentEntity, Geometry]
-    classifications: List[ClassificationAnnotation] = []
+    classifications: Optional[List[ClassificationAnnotation]] = []

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

@@ -1,16 +1,18 @@
 import abc
 from uuid import UUID, uuid4
 from typing import Any, Dict, Optional
-from labelbox import pydantic_compat
 
 from .feature import FeatureSchema
+from pydantic import PrivateAttr, ConfigDict
 
 
 class BaseAnnotation(FeatureSchema, abc.ABC):
     """ Base annotation class. Shouldn't be directly instantiated
     """
-    _uuid: Optional[UUID] = pydantic_compat.PrivateAttr()
+    _uuid: Optional[UUID] = PrivateAttr()
     extra: Dict[str, Any] = {}
+    
+    model_config = ConfigDict(extra="allow")
 
     def __init__(self, **data):
         super().__init__(**data)

libs/labelbox/src/labelbox/data/annotation_types/classification/__init__.py

@@ -1,2 +1,2 @@
-from .classification import (Checklist, ClassificationAnswer, Dropdown, Radio,
+from .classification import (Checklist, ClassificationAnswer, Radio,
                              Text)

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

@@ -1,28 +1,12 @@
 from typing import Any, Dict, List, Union, Optional
-import warnings
 from labelbox.data.annotation_types.base_annotation import BaseAnnotation
 
 from labelbox.data.mixins import ConfidenceMixin, CustomMetricsMixin
 
-try:
-    from typing import Literal
-except:
-    from typing_extensions import Literal
-
-from labelbox import pydantic_compat
+from pydantic import BaseModel
 from ..feature import FeatureSchema
 
 
-# TODO: Replace when pydantic adds support for unions that don't coerce types
-class _TempName(ConfidenceMixin, pydantic_compat.BaseModel):
-    name: str
-
-    def dict(self, *args, **kwargs):
-        res = super().dict(*args, **kwargs)
-        res.pop('name')
-        return res
-
-
 class ClassificationAnswer(FeatureSchema, ConfidenceMixin, CustomMetricsMixin):
     """
     - Represents a classification option.
@@ -36,18 +20,10 @@ class ClassificationAnswer(FeatureSchema, ConfidenceMixin, CustomMetricsMixin):
     """
     extra: Dict[str, Any] = {}
     keyframe: Optional[bool] = None
-    classifications: List['ClassificationAnnotation'] = []
+    classifications: Optional[List['ClassificationAnnotation']] = None
 
-    def dict(self, *args, **kwargs) -> Dict[str, str]:
-        res = super().dict(*args, **kwargs)
-        if res['keyframe'] is None:
-            res.pop('keyframe')
-        if res['classifications'] == []:
-            res.pop('classifications')
-        return res
 
-
-class Radio(ConfidenceMixin, CustomMetricsMixin, pydantic_compat.BaseModel):
+class Radio(ConfidenceMixin, CustomMetricsMixin, BaseModel):
     """ A classification with only one selected option allowed
 
     >>> Radio(answer = ClassificationAnswer(name = "dog"))
@@ -56,17 +32,16 @@ class Radio(ConfidenceMixin, CustomMetricsMixin, pydantic_compat.BaseModel):
     answer: ClassificationAnswer
 
 
-class Checklist(_TempName):
+class Checklist(ConfidenceMixin, BaseModel):
     """ A classification with many selected options allowed
 
     >>> Checklist(answer = [ClassificationAnswer(name = "cloudy")])
 
     """
-    name: Literal["checklist"] = "checklist"
     answer: List[ClassificationAnswer]
 
 
-class Text(ConfidenceMixin, CustomMetricsMixin, pydantic_compat.BaseModel):
+class Text(ConfidenceMixin, CustomMetricsMixin, BaseModel):
     """ Free form text
 
     >>> Text(answer = "some text answer")
@@ -75,24 +50,6 @@ class Text(ConfidenceMixin, CustomMetricsMixin, pydantic_compat.BaseModel):
     answer: str
 
 
-class Dropdown(_TempName):
-    """
-    - A classification with many selected options allowed .
-    - This is not currently compatible with MAL.
-
-    Deprecation Notice: Dropdown classification is deprecated and will be
-        removed in a future release. Dropdown will also
-        no longer be able to be created in the Editor on 3/31/2022.    
-    """
-    name: Literal["dropdown"] = "dropdown"
-    answer: List[ClassificationAnswer]
-
-    def __init__(self, **data: Any):
-        super().__init__(**data)
-        warnings.warn("Dropdown classification is deprecated and will be "
-                      "removed in a future release")
-
-
 class ClassificationAnnotation(BaseAnnotation, ConfidenceMixin,
                                CustomMetricsMixin):
     """Classification annotations (non localized)
@@ -106,12 +63,9 @@ class ClassificationAnnotation(BaseAnnotation, ConfidenceMixin,
         name (Optional[str])
         classifications (Optional[List[ClassificationAnnotation]]): Optional sub classification of the annotation
         feature_schema_id (Optional[Cuid])
-        value (Union[Text, Checklist, Radio, Dropdown])
+        value (Union[Text, Checklist, Radio])
         extra (Dict[str, Any])
      """
 
-    value: Union[Text, Checklist, Radio, Dropdown]
+    value: Union[Text, Checklist, Radio]
     message_id: Optional[str] = None
-
-
-ClassificationAnswer.update_forward_refs()

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

@@ -1,10 +1,10 @@
 from abc import ABC
 from typing import Optional, Dict, List, Any
 
-from labelbox import pydantic_compat
+from pydantic import BaseModel
 
 
-class BaseData(pydantic_compat.BaseModel, ABC):
+class BaseData(BaseModel, ABC):
     """
     Base class for objects representing data.
     This class shouldn't directly be used

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

@@ -3,5 +3,5 @@
 from .base_data import BaseData
 
 
-class ConversationData(BaseData):
-    class_name: Literal["ConversationData"] = "ConversationData"
+class ConversationData(BaseData, _NoCoercionMixin):
+    class_name: Literal["ConversationData"] = "ConversationData"

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

@@ -1,8 +1,8 @@
 from typing import Callable, Literal, Optional
 
-from labelbox import pydantic_compat
 from labelbox.data.annotation_types.data.base_data import BaseData
 from labelbox.utils import _NoCoercionMixin
+from pydantic import model_validator
 
 
 class GenericDataRowData(BaseData, _NoCoercionMixin):
@@ -14,7 +14,8 @@ class GenericDataRowData(BaseData, _NoCoercionMixin):
     def create_url(self, signer: Callable[[bytes], str]) -> Optional[str]:
         return self.url
 
-    @pydantic_compat.root_validator(pre=True)
+    @model_validator(mode="before")
+    @classmethod
     def validate_one_datarow_key_present(cls, data):
         keys = ['external_id', 'global_key', 'uid']
         count = sum([key in data for key in keys])

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

@@ -9,19 +9,22 @@
 import requests
 import numpy as np
 
-from labelbox import pydantic_compat
+from pydantic import BaseModel, model_validator, ConfigDict
 from labelbox.exceptions import InternalServerError
 from .base_data import BaseData
 from ..types import TypedArray
 
 
-class RasterData(pydantic_compat.BaseModel, ABC):
+class RasterData(BaseModel, ABC):
     """Represents an image or segmentation mask.
     """
     im_bytes: Optional[bytes] = None
     file_path: Optional[str] = None
     url: Optional[str] = None
+    uid: Optional[str] = None
+    global_key: Optional[str] = None
     arr: Optional[TypedArray[Literal['uint8']]] = None
+    model_config = ConfigDict(extra="forbid", copy_on_model_validation="none")
 
     @classmethod
     def from_2D_arr(cls, arr: Union[TypedArray[Literal['uint8']],
@@ -155,14 +158,14 @@ def create_url(self, signer: Callable[[bytes], str]) -> str:
                 "One of url, im_bytes, file_path, arr must not be None.")
         return self.url
 
-    @pydantic_compat.root_validator()
-    def validate_args(cls, values):
-        file_path = values.get("file_path")
-        im_bytes = values.get("im_bytes")
-        url = values.get("url")
-        arr = values.get("arr")
-        uid = values.get('uid')
-        global_key = values.get('global_key')
+    @model_validator(mode="after")
+    def validate_args(self, values):
+        file_path = self.file_path
+        im_bytes = self.im_bytes
+        url = self.url
+        arr = self.arr
+        uid = self.uid
+        global_key = self.global_key
         if uid == file_path == im_bytes == url == global_key == None and arr is None:
             raise ValueError(
                 "One of `file_path`, `im_bytes`, `url`, `uid`, `global_key` or `arr` required."
@@ -175,8 +178,8 @@ def validate_args(cls, values):
             elif len(arr.shape) != 3:
                 raise ValueError(
                     "unsupported image format. Must be 3D ([H,W,C])."
-                    f"Use {cls.__name__}.from_2D_arr to construct from 2D")
-        return values
+                    f"Use {self.__name__}.from_2D_arr to construct from 2D")
+        return self
 
     def __repr__(self) -> str:
         symbol_or_none = lambda data: '...' if data is not None else None
@@ -185,12 +188,6 @@ def __repr__(self) -> str:
                f"url={self.url}," \
                f"arr={symbol_or_none(self.arr)})"
 
-    class Config:
-        # Required for sharing references
-        copy_on_model_validation = 'none'
-        # Required for discriminating between data types
-        extra = 'forbid'
-
 
 class MaskData(RasterData):
     """Used to represent a segmentation Mask

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

@@ -4,7 +4,7 @@
 from requests.exceptions import ConnectTimeout
 from google.api_core import retry
 
-from labelbox import pydantic_compat
+from pydantic import ConfigDict, model_validator
 from labelbox.exceptions import InternalServerError
 from labelbox.typing_imports import Literal
 from labelbox.utils import _NoCoercionMixin
@@ -26,6 +26,7 @@ class TextData(BaseData, _NoCoercionMixin):
     file_path: Optional[str] = None
     text: Optional[str] = None
     url: Optional[str] = None
+    model_config = ConfigDict(extra="forbid")
 
     @property
     def value(self) -> str:
@@ -64,7 +65,7 @@ def fetch_remote(self) -> str:
         """
         response = requests.get(self.url)
         if response.status_code in [500, 502, 503, 504]:
-            raise labelbox.exceptions.InternalServerError(response.text)
+            raise InternalServerError(response.text)
         response.raise_for_status()
         return response.text
 
@@ -90,24 +91,20 @@ def create_url(self, signer: Callable[[bytes], str]) -> None:
                 "One of url, im_bytes, file_path, numpy must not be None.")
         return self.url
 
-    @pydantic_compat.root_validator
-    def validate_date(cls, values):
-        file_path = values.get("file_path")
-        text = values.get("text")
-        url = values.get("url")
-        uid = values.get('uid')
-        global_key = values.get('global_key')
+    @model_validator(mode="after")
+    def validate_date(self, values):
+        file_path = self.file_path
+        text = self.text
+        url = self.url
+        uid = self.uid
+        global_key = self.global_key
         if uid == file_path == text == url == global_key == None:
             raise ValueError(
                 "One of `file_path`, `text`, `uid`, `global_key` or `url` required."
             )
-        return values
+        return self
 
     def __repr__(self) -> str:
         return  f"TextData(file_path={self.file_path}," \
                 f"text={self.text[:30] + '...' if self.text is not None else None}," \
                 f"url={self.url})"
-
-    class config:
-        # Required for discriminating between data types
-        extra = 'forbid'