Dealing with BaseModel

Author: Val Brodsky

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

@@ -1,7 +1,7 @@
 import abc
 from uuid import UUID, uuid4
 from typing import Any, Dict, Optional
-from labelbox import pydantic_compat
+from pydantic import BaseModel
 
 from .feature import FeatureSchema
 

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

@@ -9,12 +9,12 @@
 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):
+class _TempName(ConfidenceMixin, BaseModel):
     name: str
 
     def dict(self, *args, **kwargs):
@@ -47,7 +47,7 @@ def dict(self, *args, **kwargs) -> Dict[str, str]:
         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"))
@@ -66,7 +66,7 @@ class Checklist(_TempName):
     answer: List[ClassificationAnswer]
 
 
-class Text(ConfidenceMixin, CustomMetricsMixin, pydantic_compat.BaseModel):
+class Text(ConfidenceMixin, CustomMetricsMixin, BaseModel):
     """ Free form text
 
     >>> Text(answer = "some text answer")

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/generic_data_row_data.py

@@ -1,6 +1,6 @@
 from typing import Callable, Literal, Optional
 
-from labelbox import pydantic_compat
+from pydantic import BaseModel, model_validator
 from labelbox.data.annotation_types.data.base_data import BaseData
 from labelbox.utils import _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(pre=True)
     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,13 +9,13 @@
 import requests
 import numpy as np
 
-from labelbox import pydantic_compat
+from pydantic import BaseModel, model_validator
 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
@@ -155,7 +155,8 @@ 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()
+    @model_validator(mode='before')
+    @classmethod()
     def validate_args(cls, values):
         file_path = values.get("file_path")
         im_bytes = values.get("im_bytes")

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 BaseModel, model_validator
 from labelbox.exceptions import InternalServerError
 from labelbox.typing_imports import Literal
 from labelbox.utils import _NoCoercionMixin
@@ -90,7 +90,8 @@ 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
+    @model_validator(mode='before')
+    @classmethod
     def validate_date(cls, values):
         file_path = values.get("file_path")
         text = values.get("text")

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

@@ -12,7 +12,7 @@
 from PIL import Image
 from pyproj import Transformer
 from pygeotile.point import Point as PygeoPoint
-from labelbox import pydantic_compat
+from pydantic import BaseModel, model_validator
 
 from labelbox.data.annotation_types import Rectangle, Point, Line, Polygon
 from .base_data import BaseData
@@ -40,7 +40,7 @@ class EPSG(Enum):
     EPSG3857 = 3857
 
 
-class TiledBounds(pydantic_compat.BaseModel):
+class TiledBounds(BaseModel):
     """ Bounds for a tiled image asset related to the relevant epsg.
 
     Bounds should be Point objects.
@@ -66,7 +66,8 @@ def validate_bounds_not_equal(cls, bounds):
         return bounds
 
     #validate bounds are within lat,lng range if they are EPSG4326
-    @pydantic_compat.root_validator
+    @model_validator(mode='before')
+    @classmethod
     def validate_bounds_lat_lng(cls, values):
         epsg = values.get('epsg')
         bounds = values.get('bounds')
@@ -82,7 +83,7 @@ def validate_bounds_lat_lng(cls, values):
         return values
 
 
-class TileLayer(pydantic_compat.BaseModel):
+class TileLayer(BaseModel):
     """ Url that contains the tile layer. Must be in the format:
 
     https://c.tile.openstreetmap.org/{z}/{x}/{y}.png
@@ -352,7 +353,7 @@ def validate_zoom_levels(cls, zoom_levels):
         return zoom_levels
 
 
-class EPSGTransformer(pydantic_compat.BaseModel):
+class EPSGTransformer(BaseModel):
     """Transformer class between different EPSG's. Useful when wanting to project
     in different formats.
     """

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

@@ -12,7 +12,7 @@
 from .base_data import BaseData
 from ..types import TypedArray
 
-from labelbox import pydantic_compat
+from pydantic import BaseModel, model_validator
 
 logger = logging.getLogger(__name__)
 
@@ -148,7 +148,8 @@ def frames_to_video(self,
         out.release()
         return file_path
 
-    @pydantic_compat.root_validator
+    @model_validator(mode='before')
+    @classmethod
     def validate_data(cls, values):
         file_path = values.get("file_path")
         url = values.get("url")

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

@@ -1,11 +1,11 @@
 from typing import Optional
 
-from labelbox import pydantic_compat
+from pydantic import BaseModel, model_validator
 
 from .types import Cuid
 
 
-class FeatureSchema(pydantic_compat.BaseModel):
+class FeatureSchema(BaseModel):
     """
     Class that represents a feature schema.
     Could be a annotation, a subclass, or an option.
@@ -14,7 +14,8 @@ class FeatureSchema(pydantic_compat.BaseModel):
     name: Optional[str] = None
     feature_schema_id: Optional[Cuid] = None
 
-    @pydantic_compat.root_validator
+    @model_validator(mode='before')
+    @classmethod
     def must_set_one(cls, values):
         if values['feature_schema_id'] is None and values['name'] is None:
             raise ValueError(

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

@@ -3,12 +3,12 @@
 
 import geojson
 import numpy as np
-from labelbox import pydantic_compat
+from pydantic import BaseModel
 
 from shapely import geometry as geom
 
 
-class Geometry(pydantic_compat.BaseModel, ABC):
+class Geometry(BaseModel, ABC):
     """Abstract base class for geometry objects
     """
     extra: Dict[str, Any] = {}