docarray · anna-charlotte · Apr 6, 2023 · Apr 5, 2023 · Apr 6, 2023 · Apr 6, 2023
diff --git a/docarray/typing/bytes/image_bytes.py b/docarray/typing/bytes/image_bytes.py
@@ -7,7 +7,7 @@
 
 from docarray.typing.abstract_type import AbstractType
 from docarray.typing.proto_register import _register_proto
-from docarray.typing.tensor.image import ImageNdArray
+from docarray.typing.tensor.image.image_ndarray import ImageNdArray
 from docarray.utils._internal.misc import import_library
 
 if TYPE_CHECKING:

diff --git a/docarray/typing/tensor/audio/abstract_audio_tensor.py b/docarray/typing/tensor/audio/abstract_audio_tensor.py
@@ -5,19 +5,24 @@
 from docarray.typing.tensor.abstract_tensor import AbstractTensor
 from docarray.utils._internal.misc import import_library, is_notebook
 
+if TYPE_CHECKING:
+    from docarray.typing.bytes.audio_bytes import AudioBytes
+
 T = TypeVar('T', bound='AbstractAudioTensor')
 
 MAX_INT_16 = 2**15
 
 
 class AbstractAudioTensor(AbstractTensor, ABC):
-    def to_bytes(self):
+    def to_bytes(self) -> 'AudioBytes':
         """
-        Convert audio tensor to bytes.
+        Convert audio tensor to AudioBytes.
         """
+        from docarray.typing.bytes.audio_bytes import AudioBytes
+
         tensor = self.get_comp_backend().to_numpy(self)
         tensor = (tensor * MAX_INT_16).astype('<h')
-        return tensor.tobytes()
+        return AudioBytes(tensor.tobytes())
 
     def save(
         self: 'T',

diff --git a/docarray/typing/tensor/audio/audio_ndarray.py b/docarray/typing/tensor/audio/audio_ndarray.py
@@ -16,15 +16,15 @@ class AudioNdArray(AbstractAudioTensor, NdArray):
     from typing import Optional
 
     from docarray import BaseDoc
-    from docarray.typing import AudioNdArray, AudioUrl
+    from docarray.typing import AudioBytes, AudioNdArray, AudioUrl
     import numpy as np
 
 
     class MyAudioDoc(BaseDoc):
         title: str
         audio_tensor: Optional[AudioNdArray]
         url: Optional[AudioUrl]
-        bytes_: Optional[bytes]
+        bytes_: Optional[AudioBytes]
 
 
     # from tensor

diff --git a/docarray/typing/tensor/audio/audio_tensorflow_tensor.py b/docarray/typing/tensor/audio/audio_tensorflow_tensor.py
@@ -26,14 +26,14 @@ class AudioTensorFlowTensor(
     from pydantic import parse_obj_as
 
     from docarray import BaseDoc
-    from docarray.typing import AudioTensorFlowTensor, AudioUrl
+    from docarray.typing import AudioBytes, AudioTensorFlowTensor, AudioUrl
 
 
     class MyAudioDoc(BaseDoc):
         title: str
         audio_tensor: Optional[AudioTensorFlowTensor]
         url: Optional[AudioUrl]
-        bytes_: Optional[bytes]
+        bytes_: Optional[AudioBytes]
 
 
     doc_1 = MyAudioDoc(

diff --git a/docarray/typing/tensor/audio/audio_torch_tensor.py b/docarray/typing/tensor/audio/audio_torch_tensor.py
@@ -18,14 +18,14 @@ class AudioTorchTensor(AbstractAudioTensor, TorchTensor, metaclass=metaTorchAndN
     import torch
 
     from docarray import BaseDoc
-    from docarray.typing import AudioTorchTensor, AudioUrl
+    from docarray.typing import AudioBytes, AudioTorchTensor, AudioUrl
 
 
     class MyAudioDoc(BaseDoc):
         title: str
         audio_tensor: Optional[AudioTorchTensor]
         url: Optional[AudioUrl]
-        bytes_: Optional[bytes]
+        bytes_: Optional[AudioBytes]
 
 
     doc_1 = MyAudioDoc(

diff --git a/docarray/typing/tensor/image/abstract_image_tensor.py b/docarray/typing/tensor/image/abstract_image_tensor.py
@@ -3,18 +3,22 @@
 from abc import ABC
 
 import numpy as np
+from typing_extensions import TYPE_CHECKING
 
 from docarray.typing.tensor.abstract_tensor import AbstractTensor
 from docarray.utils._internal.misc import import_library, is_notebook
 
+if TYPE_CHECKING:
+    from docarray.typing.bytes.image_bytes import ImageBytes
+
 
 class AbstractImageTensor(AbstractTensor, ABC):
-    def to_bytes(self, format: str = 'PNG') -> bytes:
+    def to_bytes(self, format: str = 'PNG') -> 'ImageBytes':
         """
-        Convert image tensor to bytes.
+        Convert image tensor to ImageBytes.
 
         :param format: the image format use to store the image, can be 'PNG' , 'JPG' ...
-        :return: bytes
+        :return: an ImageBytes object
         """
         PIL = import_library('PIL', raise_error=True)  # noqa: F841
         from PIL import Image as PILImage
@@ -31,7 +35,9 @@ def to_bytes(self, format: str = 'PNG') -> bytes:
             pil_image.save(buffer, format=format)
             img_byte_arr = buffer.getvalue()
 
-        return img_byte_arr
+        from docarray.typing.bytes.image_bytes import ImageBytes
+
+        return ImageBytes(img_byte_arr)
 
     def save(self, file_path: str) -> None:
         """

diff --git a/docarray/typing/tensor/image/image_ndarray.py b/docarray/typing/tensor/image/image_ndarray.py
@@ -20,14 +20,14 @@ class ImageNdArray(AbstractImageTensor, NdArray):
     from typing import Optional
 
     from docarray import BaseDoc
-    from docarray.typing import ImageNdArray, ImageUrl
+    from docarray.typing import ImageBytes, ImageNdArray, ImageUrl
 
 
     class MyImageDoc(BaseDoc):
         title: str
         tensor: Optional[ImageNdArray]
         url: Optional[ImageUrl]
-        bytes: Optional[bytes]
+        bytes: Optional[ImageBytes]
 
 
     # from url

diff --git a/docarray/typing/tensor/image/image_tensorflow_tensor.py b/docarray/typing/tensor/image/image_tensorflow_tensor.py
@@ -24,14 +24,14 @@ class ImageTensorFlowTensor(
     from typing import Optional
 
     from docarray import BaseDoc
-    from docarray.typing import ImageTensorFlowTensor, ImageUrl
+    from docarray.typing import ImageBytes, ImageTensorFlowTensor, ImageUrl
 
 
     class MyImageDoc(BaseDoc):
         title: str
         tensor: Optional[ImageTensorFlowTensor]
         url: Optional[ImageUrl]
-        bytes: Optional[bytes]
+        bytes: Optional[ImageBytes]
 
 
     doc = MyImageDoc(

diff --git a/docarray/typing/tensor/image/image_torch_tensor.py b/docarray/typing/tensor/image/image_torch_tensor.py
@@ -22,14 +22,14 @@ class ImageTorchTensor(AbstractImageTensor, TorchTensor, metaclass=metaTorchAndN
     from typing import Optional
 
     from docarray import BaseDoc
-    from docarray.typing import ImageTorchTensor, ImageUrl
+    from docarray.typing import ImageBytes, ImageTorchTensor, ImageUrl
 
 
     class MyImageDoc(BaseDoc):
         title: str
         tensor: Optional[ImageTorchTensor]
         url: Optional[ImageUrl]
-        bytes: Optional[bytes]
+        bytes: Optional[ImageBytes]
 
 
     doc = MyImageDoc(

diff --git a/docarray/typing/tensor/video/video_tensor_mixin.py b/docarray/typing/tensor/video/video_tensor_mixin.py
@@ -9,6 +9,9 @@
 from docarray.typing.tensor.audio.audio_tensor import AudioTensor
 from docarray.utils._internal.misc import import_library, is_notebook
 
+if TYPE_CHECKING:
+    from docarray.typing.bytes.video_bytes import VideoBytes
+
 T = TypeVar('T', bound='VideoTensorMixin')
 
 
@@ -130,9 +133,9 @@ def to_bytes(
         audio_frame_rate: int = 48000,
         audio_codec: str = 'aac',
         audio_format: str = 'fltp',
-    ) -> bytes:
+    ) -> 'VideoBytes':
         """
-        Convert video tensor to bytes.
+        Convert video tensor to VideoBytes.
 
         :param audio_tensor: AudioTensor containing the video's soundtrack.
         :param video_frame_rate: video frames per second.
@@ -142,8 +145,10 @@ def to_bytes(
         :param audio_format: the name of one of the audio formats supported by PyAV,
             such as 'flt', 'fltp', 's16' or 's16p'.
 
-        :return: bytes
+        :return: a VideoBytes object
         """
+        from docarray.typing.bytes.video_bytes import VideoBytes
+
         bytes = BytesIO()
         self.save(
             file_path=bytes,
@@ -154,7 +159,7 @@ def to_bytes(
             audio_codec=audio_codec,
             audio_format=audio_format,
         )
-        return bytes.getvalue()
+        return VideoBytes(bytes.getvalue())
 
     def display(self, audio: Optional[AudioTensor] = None) -> None:
         """

diff --git a/tests/integrations/predefined_document/test_image.py b/tests/integrations/predefined_document/test_image.py
@@ -84,4 +84,5 @@ def test_byte_from_tensor():
     img.bytes_ = img.tensor.to_bytes()
 
     assert isinstance(img.bytes_, bytes)
+    assert isinstance(img.bytes_, ImageBytes)
     assert len(img.bytes_) > 0
diff --git a/tests/units/typing/tensor/test_audio_tensor.py b/tests/units/typing/tensor/test_audio_tensor.py
@@ -6,6 +6,7 @@
 from pydantic import parse_obj_as
 
 from docarray import BaseDoc
+from docarray.typing.bytes.audio_bytes import AudioBytes
 from docarray.typing.tensor.audio.audio_ndarray import AudioNdArray
 from docarray.typing.tensor.audio.audio_torch_tensor import AudioTorchTensor
 from docarray.utils._internal.misc import is_tf_available
@@ -120,3 +121,16 @@ def test_save_audio_tensorflow_tensor_to_wav_file(tmpdir):
     audio_tensor = parse_obj_as(AudioTensorFlowTensor, tf.zeros((1000, 2)))
     audio_tensor.save(tmp_file)
     assert os.path.isfile(tmp_file)
+
+
+@pytest.mark.parametrize(
+    'audio_tensor',
+    [
+        parse_obj_as(AudioTorchTensor, torch.zeros(1000, 2)),
+        parse_obj_as(AudioNdArray, np.zeros((1000, 2))),
+    ],
+)
+def test_save_audio_tensor_to_bytes(audio_tensor):
+    b = audio_tensor.to_bytes()
+    isinstance(b, bytes)
+    isinstance(b, AudioBytes)
diff --git a/tests/units/typing/tensor/test_image_tensor.py b/tests/units/typing/tensor/test_image_tensor.py
@@ -5,7 +5,7 @@
 import torch
 from pydantic import parse_obj_as
 
-from docarray.typing import ImageNdArray, ImageTorchTensor
+from docarray.typing import ImageBytes, ImageNdArray, ImageTorchTensor
 from docarray.utils._internal.misc import is_tf_available
 
 tf_available = is_tf_available()
@@ -35,3 +35,16 @@ def test_save_image_tensorflow_tensor_to_file(tmpdir):
     image_tensor = parse_obj_as(ImageTensorFlowTensor, tf.zeros((224, 224, 3)))
     image_tensor.save(tmp_file)
     assert os.path.isfile(tmp_file)
+
+
+@pytest.mark.parametrize(
+    'image_tensor',
+    [
+        parse_obj_as(ImageTorchTensor, torch.zeros(224, 224, 3)),
+        parse_obj_as(ImageNdArray, np.zeros((224, 224, 3))),
+    ],
+)
+def test_save_image_tensor_to_bytes(image_tensor):
+    b = image_tensor.to_bytes()
+    isinstance(b, bytes)
+    isinstance(b, ImageBytes)
diff --git a/tests/units/typing/tensor/test_video_tensor.py b/tests/units/typing/tensor/test_video_tensor.py
@@ -9,6 +9,7 @@
 from docarray.typing import (
     AudioNdArray,
     AudioTorchTensor,
+    VideoBytes,
     VideoNdArray,
     VideoTorchTensor,
 )
@@ -137,9 +138,10 @@ def test_save_video_tensor_to_file(video_tensor, tmpdir):
         parse_obj_as(VideoNdArray, np.zeros((1, 224, 224, 3))),
     ],
 )
-def test_save_video_tensor_to_bytes(video_tensor, tmpdir):
+def test_save_video_tensor_to_bytes(video_tensor):
     b = video_tensor.to_bytes()
     isinstance(b, bytes)
+    isinstance(b, VideoBytes)
 
 
 @pytest.mark.tensorflow