Deal with TypedArray special validations

Author: Val Brodsky

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

@@ -155,29 +155,22 @@ def create_url(self, signer: Callable[[bytes], str]) -> str:
                 "One of url, im_bytes, file_path, arr must not be None.")
         return self.url
 
-    @model_validator(mode='before')
-    @classmethod
-    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')
-        if uid == file_path == im_bytes == url == global_key == None and arr is None:
+    @model_validator(mode='after')
+    def validate_args(self):
+        if self.uid == self.file_path == self.im_bytes == self.url == self.global_key == None and self.arr is None:
             raise ValueError(
                 "One of `file_path`, `im_bytes`, `url`, `uid`, `global_key` or `arr` required."
             )
-        if arr is not None:
-            if arr.dtype != np.uint8:
+        if self.arr is not None:
+            if self.arr.dtype != np.uint8:
                 raise TypeError(
                     "Numpy array representing segmentation mask must be np.uint8"
                 )
-            elif len(arr.shape) != 3:
+            elif len(self.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
+        return self
 
     def __repr__(self) -> str:
         symbol_or_none = lambda data: '...' if data is not None else None

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

@@ -148,20 +148,13 @@ def frames_to_video(self,
         out.release()
         return file_path
 
-    @model_validator(mode='before')
-    @classmethod
-    def validate_data(cls, values):
-        file_path = values.get("file_path")
-        url = values.get("url")
-        frames = values.get("frames")
-        uid = values.get("uid")
-        global_key = values.get("global_key")
-
-        if uid == file_path == frames == url == global_key == None:
+    @model_validator(mode='after')
+    def validate_data(self):
+        if self.uid == self.file_path == self.frames == self.url == self.global_key == None:
             raise ValueError(
                 "One of `file_path`, `frames`, `uid`, `global_key` or `url` required."
             )
-        return values
+        return self
 
     def __repr__(self) -> str:
         return  f"VideoData(file_path={self.file_path}," \

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

@@ -1,11 +1,12 @@
 import sys
-from typing import Generic, TypeVar, Any
+from typing import Generic, TypeVar, Any, Type
 
 from labelbox.typing_imports import Annotated
 from packaging import version
 import numpy as np
 
-from pydantic import ValidationInfo, Field
+from pydantic import Field, GetCoreSchemaHandler, TypeAdapter
+from pydantic_core import core_schema
 
 Cuid = Annotated[str, Field(min_length=25, max_length=25)]
 
@@ -15,18 +16,28 @@
 
 class _TypedArray(np.ndarray, Generic[DType, DShape]):
 
+    # @classmethod
+    # def __get_validators__(cls):
+    #     yield cls.validate
+
     @classmethod
-    def __get_validators__(cls):
-        yield cls.validate
+    def __get_pydantic_core_schema__(
+            cls, source_type: Type[Any],
+            handler: GetCoreSchemaHandler) -> core_schema.CoreSchema:
+
+        # assert source is CompressedString
+        return core_schema.with_info_after_validator_function(
+            function=cls.validate,
+            schema=core_schema.any_schema(),
+            field_name=source_type.__args__[-1].__args__[0])
 
     @classmethod
-    def validate(cls, val, field_info: ValidationInfo):
+    def validate(cls, val, info):
         if not isinstance(val, np.ndarray):
             raise TypeError(f"Expected numpy array. Found {type(val)}")
 
-        actual_dtype = cls.model_fields[field_info.name].type_.__args__[0]
-
-        if val.dtype != actual_dtype:
+        actual_type = info.field_name
+        if str(val.dtype) != actual_type:
             raise TypeError(
                 f"Expected numpy array have type {actual_dtype}. Found {val.dtype}"
             )