fix: fix some stuff

Signed-off-by: samsja <[email protected]>
docarray · samsja · Sep 10, 2023 · Jun 15, 2023 · Jun 15, 2023 · Jun 15, 2023
commit c20f49baf01ef5649471145a8841a1e1307a6b07
diff --git a/docarray/typing/bytes/audio_bytes.py b/docarray/typing/bytes/audio_bytes.py
@@ -1,44 +1,23 @@
 import io
-from typing import TYPE_CHECKING, Any, Tuple, Type, TypeVar
+from typing import Tuple, TypeVar
 
 import numpy as np
 from pydantic import parse_obj_as
 
-from docarray.typing.abstract_type import AbstractType
+from docarray.typing.bytes.base_bytes import BaseBytes
 from docarray.typing.proto_register import _register_proto
 from docarray.typing.tensor.audio import AudioNdArray
 from docarray.utils._internal.misc import import_library
-from docarray.utils._internal.pydantic import bytes_validator
-
-if TYPE_CHECKING:
-    from docarray.proto import NodeProto
 
 T = TypeVar('T', bound='AudioBytes')
 
 
 @_register_proto(proto_type_name='audio_bytes')
-class AudioBytes(bytes, AbstractType):
+class AudioBytes(BaseBytes):
     """
     Bytes that store an audio and that can be load into an Audio tensor
     """
 
-    @classmethod
-    def _docarray_validate(
-        cls: Type[T],
-        value: Any,
-    ) -> T:
-        value = bytes_validator(value)
-        return cls(value)
-
-    @classmethod
-    def from_protobuf(cls: Type[T], pb_msg: T) -> T:
-        return parse_obj_as(cls, pb_msg)
-
-    def _to_node_protobuf(self: T) -> 'NodeProto':
-        from docarray.proto import NodeProto
-
-        return NodeProto(blob=self, type=self._proto_type_name)
-
     def load(self) -> Tuple[AudioNdArray, int]:
         """
         Load the Audio from the [`AudioBytes`][docarray.typing.AudioBytes] into an

diff --git a/docarray/typing/bytes/base_bytes.py b/docarray/typing/bytes/base_bytes.py
@@ -0,0 +1,53 @@
+from abc import abstractmethod
+from typing import TYPE_CHECKING, Any, Type, TypeVar
+
+from pydantic import parse_obj_as
+
+from docarray.typing.abstract_type import AbstractType
+from docarray.utils._internal.pydantic import bytes_validator, is_pydantic_v2
+
+if is_pydantic_v2:
+    from pydantic_core import core_schema
+
+if TYPE_CHECKING:
+    from docarray.proto import NodeProto
+
+    if is_pydantic_v2:
+        from pydantic import GetCoreSchemaHandler
+
+T = TypeVar('T', bound='BaseBytes')
+
+
+class BaseBytes(bytes, AbstractType):
+    """
+    Bytes type for docarray
+    """
+
+    @classmethod
+    def _docarray_validate(
+        cls: Type[T],
+        value: Any,
+    ) -> T:
+        value = bytes_validator(value)
+        return cls(value)
+
+    @classmethod
+    def from_protobuf(cls: Type[T], pb_msg: T) -> T:
+        return parse_obj_as(cls, pb_msg)
+
+    def _to_node_protobuf(self: T) -> 'NodeProto':
+        from docarray.proto import NodeProto
+
+        return NodeProto(blob=self, type=self._proto_type_name)
+
+    if is_pydantic_v2:
+
+        @classmethod
+        @abstractmethod
+        def __get_pydantic_core_schema__(
+            cls, _source_type: Any, _handler: 'GetCoreSchemaHandler'
+        ) -> 'core_schema.CoreSchema':
+            return core_schema.general_after_validator_function(
+                cls.validate,
+                core_schema.bytes_schema(),
+            )
diff --git a/docarray/typing/bytes/image_bytes.py b/docarray/typing/bytes/image_bytes.py
@@ -1,46 +1,27 @@
 from io import BytesIO
-from typing import TYPE_CHECKING, Any, Optional, Tuple, Type, TypeVar
+from typing import TYPE_CHECKING, Optional, Tuple, TypeVar
 
 import numpy as np
 from pydantic import parse_obj_as
 
-from docarray.typing.abstract_type import AbstractType
+from docarray.typing.bytes.base_bytes import BaseBytes
 from docarray.typing.proto_register import _register_proto
 from docarray.typing.tensor.image.image_ndarray import ImageNdArray
 from docarray.utils._internal.misc import import_library
-from docarray.utils._internal.pydantic import bytes_validator
 
 if TYPE_CHECKING:
     from PIL import Image as PILImage
 
-    from docarray.proto import NodeProto
 
 T = TypeVar('T', bound='ImageBytes')
 
 
 @_register_proto(proto_type_name='image_bytes')
-class ImageBytes(bytes, AbstractType):
+class ImageBytes(BaseBytes):
     """
     Bytes that store an image and that can be load into an image tensor
     """
 
-    @classmethod
-    def _docarray_validate(
-        cls: Type[T],
-        value: Any,
-    ) -> T:
-        value = bytes_validator(value)
-        return cls(value)
-
-    @classmethod
-    def from_protobuf(cls: Type[T], pb_msg: T) -> T:
-        return parse_obj_as(cls, pb_msg)
-
-    def _to_node_protobuf(self: T) -> 'NodeProto':
-        from docarray.proto import NodeProto
-
-        return NodeProto(blob=self, type=self._proto_type_name)
-
     def load_pil(
         self,
     ) -> 'PILImage.Image':

diff --git a/docarray/typing/bytes/video_bytes.py b/docarray/typing/bytes/video_bytes.py
@@ -1,17 +1,13 @@
 from io import BytesIO
-from typing import TYPE_CHECKING, Any, List, NamedTuple, Type, TypeVar
+from typing import TYPE_CHECKING, List, NamedTuple, TypeVar
 
 import numpy as np
 from pydantic import parse_obj_as
 
-from docarray.typing.abstract_type import AbstractType
+from docarray.typing.bytes.base_bytes import BaseBytes
 from docarray.typing.proto_register import _register_proto
 from docarray.typing.tensor import AudioNdArray, NdArray, VideoNdArray
 from docarray.utils._internal.misc import import_library
-from docarray.utils._internal.pydantic import bytes_validator
-
-if TYPE_CHECKING:
-    from docarray.proto import NodeProto
 
 T = TypeVar('T', bound='VideoBytes')
 
@@ -23,28 +19,11 @@ class VideoLoadResult(NamedTuple):
 
 
 @_register_proto(proto_type_name='video_bytes')
-class VideoBytes(bytes, AbstractType):
+class VideoBytes(BaseBytes):
     """
     Bytes that store a video and that can be load into a video tensor
     """
 
-    @classmethod
-    def _docarray_validate(
-        cls: Type[T],
-        value: Any,
-    ) -> T:
-        value = bytes_validator(value)
-        return cls(value)
-
-    @classmethod
-    def from_protobuf(cls: Type[T], pb_msg: T) -> T:
-        return parse_obj_as(cls, pb_msg)
-
-    def _to_node_protobuf(self: T) -> 'NodeProto':
-        from docarray.proto import NodeProto
-
-        return NodeProto(blob=self, type=self._proto_type_name)
-
     def load(self, **kwargs) -> VideoLoadResult:
         """
         Load the video from the bytes into a VideoLoadResult object consisting of:

diff --git a/docarray/typing/url/any_url.py b/docarray/typing/url/any_url.py
@@ -7,6 +7,7 @@
 import numpy as np
 from pydantic import AnyUrl as BaseAnyUrl
 from pydantic import errors, parse_obj_as
+from pydantic_core import core_schema
 
 from docarray.typing.abstract_type import AbstractType
 from docarray.typing.proto_register import _register_proto
@@ -16,6 +17,8 @@
     if not is_pydantic_v2:
         from pydantic import BaseConfig
         from pydantic.fields import ModelField
+    else:
+        from pydantic import GetCoreSchemaHandler
 
     from pydantic.networks import Parts
 
@@ -27,9 +30,25 @@
 if is_pydantic_v2:
 
     @_register_proto(proto_type_name='any_url')
-    class AnyUrl:
+    class AnyUrl(AbstractType):
         def __init__(self, *args, **kwargs):
-            raise NotImplementedError('AnyUrl is not supported in pydantic v2')
+            raise NotImplementedError('AnyUrl is not supported in pydantic v2 for now')
+
+        @classmethod
+        def _docarray_validate(
+            cls: Type[T],
+            value: Any,
+        ):
+            raise NotImplementedError('AnyUrl is not supported in pydantic v2 for now')
+
+        def __get_pydantic_core_schema__(
+            cls, source: type[Any], handler: Optional['GetCoreSchemaHandler'] = None
+        ) -> core_schema.CoreSchema:
+
+            return core_schema.general_after_validator_function(
+                cls._docarray_validate,
+                core_schema.str_schema(),
+            )
 
 else:
 

diff --git a/docarray/utils/_internal/pydantic.py b/docarray/utils/_internal/pydantic.py
@@ -7,6 +7,6 @@
     from pydantic.validators import bytes_validator
 
 else:
+    from pydantic.v1.validators import bytes_validator
 
-    def bytes_validator(*args, **kwargs):
-        raise NotImplementedError('bytes_validator is not implemented in pydantic v2')
+__all__ = ['is_pydantic_v2', 'bytes_validator']