docarray · anna-charlotte · Feb 9, 2023 · Jan 30, 2023 · Jan 30, 2023 · Jan 30, 2023
diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml
@@ -21,6 +21,6 @@ repos:
         exclude: ^(docarray/proto/pb/docarray_pb2.py|docarray/proto/pb/docarray_pb2.py|docs/|docarray/resources/)
 
 - repo: https://github.com/charliermarsh/ruff-pre-commit
-  rev: v0.0.165
+  rev: v0.0.243
   hooks:
     - id: ruff
diff --git a/docarray/array/array_stacked.py b/docarray/array/array_stacked.py
@@ -21,27 +21,27 @@
 from docarray.typing import NdArray
 from docarray.typing.tensor.abstract_tensor import AbstractTensor
 from docarray.utils._typing import is_tensor_union
+from docarray.utils.misc import is_tf_available, is_torch_available
 
 if TYPE_CHECKING:
     from pydantic import BaseConfig
     from pydantic.fields import ModelField
 
     from docarray.proto import DocumentArrayStackedProto
 
-try:
+torch_available = is_torch_available()
+if torch_available:
     from docarray.typing import TorchTensor
-except ImportError:
+else:
     TorchTensor = None  # type: ignore
 
-try:
+tf_available = is_tf_available()
+if tf_available:
     import tensorflow as tf  # type: ignore
 
     from docarray.typing import TensorFlowTensor
-
-    tf_available = True
-except (ImportError, TypeError):
+else:
     TensorFlowTensor = None  # type: ignore
-    tf_available = False
 
 T = TypeVar('T', bound='DocumentArrayStacked')
 IndexIterType = Union[slice, Iterable[int], Iterable[bool], None]

diff --git a/docarray/documents/audio.py b/docarray/documents/audio.py
@@ -7,12 +7,17 @@
 from docarray.typing.bytes.audio_bytes import AudioBytes
 from docarray.typing.tensor.abstract_tensor import AbstractTensor
 from docarray.typing.tensor.audio.audio_tensor import AudioTensor
-from docarray.utils.misc import is_torch_available
+from docarray.utils.misc import is_tf_available, is_torch_available
 
 torch_available = is_torch_available()
 if torch_available:
     import torch
 
+tf_available = is_tf_available()
+if tf_available:
+    import tensorflow as tf  # type: ignore
+
+
 T = TypeVar('T', bound='Audio')
 
 
@@ -102,7 +107,9 @@ def validate(
         if isinstance(value, str):
             value = cls(url=value)
         elif isinstance(value, (AbstractTensor, np.ndarray)) or (
-            torch_available and isinstance(value, torch.Tensor)
+            torch_available
+            and isinstance(value, torch.Tensor)
+            or (tf_available and isinstance(value, tf.Tensor))
         ):
             value = cls(tensor=value)
 

diff --git a/docarray/documents/image.py b/docarray/documents/image.py
@@ -6,14 +6,18 @@
 from docarray.typing import AnyEmbedding, ImageBytes, ImageUrl
 from docarray.typing.tensor.abstract_tensor import AbstractTensor
 from docarray.typing.tensor.image.image_tensor import ImageTensor
-from docarray.utils.misc import is_torch_available
+from docarray.utils.misc import is_tf_available, is_torch_available
 
 T = TypeVar('T', bound='Image')
 
 torch_available = is_torch_available()
 if torch_available:
     import torch
 
+tf_available = is_tf_available()
+if tf_available:
+    import tensorflow as tf  # type: ignore
+
 
 class Image(BaseDocument):
     """
@@ -91,8 +95,10 @@ def validate(
     ) -> T:
         if isinstance(value, str):
             value = cls(url=value)
-        elif isinstance(value, (AbstractTensor, np.ndarray)) or (
-            torch_available and isinstance(value, torch.Tensor)
+        elif (
+            isinstance(value, (AbstractTensor, np.ndarray))
+            or (torch_available and isinstance(value, torch.Tensor))
+            or (tf_available and isinstance(value, tf.Tensor))
         ):
             value = cls(tensor=value)
         elif isinstance(value, bytes):

diff --git a/docarray/documents/point_cloud.py b/docarray/documents/point_cloud.py
@@ -5,12 +5,16 @@
 from docarray.base_document import BaseDocument
 from docarray.typing import AnyEmbedding, AnyTensor, PointCloud3DUrl
 from docarray.typing.tensor.abstract_tensor import AbstractTensor
-from docarray.utils.misc import is_torch_available
+from docarray.utils.misc import is_tf_available, is_torch_available
 
 torch_available = is_torch_available()
 if torch_available:
     import torch
 
+tf_available = is_tf_available()
+if tf_available:
+    import tensorflow as tf  # type: ignore
+
 T = TypeVar('T', bound='PointCloud3D')
 
 
@@ -100,7 +104,9 @@ def validate(
         if isinstance(value, str):
             value = cls(url=value)
         elif isinstance(value, (AbstractTensor, np.ndarray)) or (
-            torch_available and isinstance(value, torch.Tensor)
+            torch_available
+            and isinstance(value, torch.Tensor)
+            or (tf_available and isinstance(value, tf.Tensor))
         ):
             value = cls(tensor=value)
 

diff --git a/docarray/documents/text.py b/docarray/documents/text.py
@@ -74,7 +74,7 @@ class MultiModalDoc(BaseDocument):
         )
         mmdoc.text_doc.text = mmdoc.text_doc.url.load()
 
-        #or
+        # or
 
         mmdoc.text_doc.bytes = mmdoc.text_doc.url.load_bytes()
 
@@ -87,13 +87,13 @@ class MultiModalDoc(BaseDocument):
 
     .. code-block:: python
 
-        from docarray.documents Text
+        from docarray.documents import Text
 
         doc = Text(text='This is the main text', url='exampleurl.com')
         doc2 = Text(text='This is the main text', url='exampleurl.com')
 
-        doc == 'This is the main text' # True
-        doc == doc2 # False, their ids are not equivalent
+        doc == 'This is the main text'  # True
+        doc == doc2  # False, their ids are not equivalent
     """
 
     text: Optional[str] = None
@@ -126,7 +126,7 @@ def __contains__(self, item: str) -> bool:
         """
         This method makes `Text` behave the same as an `str`.
 
-            .. code-block:: python
+        .. code-block:: python
 
             from docarray.documents import Text
 

diff --git a/docarray/documents/video.py b/docarray/documents/video.py
@@ -8,12 +8,18 @@
 from docarray.typing.tensor.abstract_tensor import AbstractTensor
 from docarray.typing.tensor.video.video_tensor import VideoTensor
 from docarray.typing.url.video_url import VideoUrl
-from docarray.utils.misc import is_torch_available
+from docarray.utils.misc import is_tf_available, is_torch_available
 
 torch_available = is_torch_available()
 if torch_available:
     import torch
 
+
+tf_available = is_tf_available()
+if tf_available:
+    import tensorflow as tf  # type: ignore
+
+
 T = TypeVar('T', bound='Video')
 
 
@@ -106,7 +112,9 @@ def validate(
         if isinstance(value, str):
             value = cls(url=value)
         elif isinstance(value, (AbstractTensor, np.ndarray)) or (
-            torch_available and isinstance(value, torch.Tensor)
+            torch_available
+            and isinstance(value, torch.Tensor)
+            or (tf_available and isinstance(value, tf.Tensor))
         ):
             value = cls(tensor=value)
 

diff --git a/docarray/typing/__init__.py b/docarray/typing/__init__.py
@@ -38,11 +38,10 @@
     'ImageNdArray',
 ]
 
-try:
-    import torch  # noqa: F401
-except ImportError:
-    pass
-else:
+from docarray.utils.misc import is_tf_available, is_torch_available
+
+torch_available = is_torch_available()
+if torch_available:
     from docarray.typing.tensor import TorchEmbedding, TorchTensor  # noqa: F401
     from docarray.typing.tensor.audio.audio_torch_tensor import AudioTorchTensor  # noqa
     from docarray.typing.tensor.image import ImageTorchTensor  # noqa:  F401
@@ -58,11 +57,20 @@
         ]
     )
 
-try:
-    import tensorflow as tf  # type: ignore # noqa: F401
-except (ImportError, TypeError):
-    pass
-else:
-    from docarray.typing.tensor import TensorFlowTensor  # noqa: F401
+tf_available = is_tf_available()
+if tf_available:
+    from docarray.typing.tensor import TensorFlowTensor
+    from docarray.typing.tensor.audio import AudioTensorFlowTensor  # noqa: F401
    from docarray.typing.tensor.embedding import TensorFlowEmbedding  # noqa: F401
+    from docarray.typing.tensor.image import ImageTensorFlowTensor  # noqa: F401
+    from docarray.typing.tensor.video import VideoTensorFlowTensor  # noqa
 
-    __all__.extend(['TensorFlowTensor'])
+    __all__.extend(
+        [
+            'TensorFlowTensor',
+            'TensorFlowEmbedding',
+            'AudioTensorFlowTensor',
+            'ImageTensorFlowTensor',
+            'VideoTensorFlowTensor',
+        ]
+    )
diff --git a/docarray/typing/tensor/__init__.py b/docarray/typing/tensor/__init__.py
@@ -13,22 +13,23 @@
     'TensorFlowTensor',
 ]
 
-try:
-    import torch  # noqa: F401
-except ImportError:
-    pass
-else:
+from docarray.utils.misc import is_tf_available, is_torch_available
+
+torch_available = is_torch_available()
+if torch_available:
     from docarray.typing.tensor.embedding import TorchEmbedding  # noqa: F401
     from docarray.typing.tensor.image import ImageTorchTensor  # noqa: F401
     from docarray.typing.tensor.torch_tensor import TorchTensor  # noqa: F401
 
     __all__.extend(['TorchEmbedding', 'TorchTensor', 'ImageTorchTensor'])
 
-try:
-    import tensorflow as tf  # type: ignore # noqa: F401
-except (ImportError, TypeError):
-    pass
-else:
+    torch_available = is_torch_available()
+
+
+tf_available = is_tf_available()
+if tf_available:
+    from docarray.typing.tensor.embedding import TensorFlowEmbedding  # noqa: F401
+    from docarray.typing.tensor.image import ImageTensorFlowTensor  # noqa: F401
     from docarray.typing.tensor.tensorflow_tensor import TensorFlowTensor  # noqa: F401
 
-    __all__.extend(['TensorFlowTensor'])
+    __all__.extend(['TensorFlowEmbedding', 'TensorFlowTensor', 'ImageTensorFlowTensor'])
diff --git a/docarray/typing/tensor/audio/__init__.py b/docarray/typing/tensor/audio/__init__.py
@@ -2,11 +2,19 @@
 
 __all__ = ['AudioNdArray']
 
-try:
-    import torch  # noqa: F401
-except ImportError:
-    pass
-else:
+from docarray.utils.misc import is_tf_available, is_torch_available
+
+torch_available = is_torch_available()
+if torch_available:
     from docarray.typing.tensor.audio.audio_torch_tensor import AudioTorchTensor  # noqa
 
     __all__.extend(['AudioTorchTensor'])
+
+
+tf_available = is_tf_available()
+if tf_available:
+    from docarray.typing.tensor.audio.audio_tensorflow_tensor import (  # noqa
+        AudioTensorFlowTensor,
+    )
+
+    __all__.extend(['AudioTensorFlowTensor'])
diff --git a/docarray/typing/tensor/audio/audio_ndarray.py b/docarray/typing/tensor/audio/audio_ndarray.py
@@ -16,14 +16,12 @@ class AudioNdArray(AbstractAudioTensor, NdArray):
 
         from typing import Optional
 
-        from pydantic import parse_obj_as
-
-        from docarray import Document
+        from docarray import BaseDocument
         from docarray.typing import AudioNdArray, AudioUrl
         import numpy as np
 
 
-        class MyAudioDoc(Document):
+        class MyAudioDoc(BaseDocument):
             title: str
             audio_tensor: Optional[AudioNdArray]
             url: Optional[AudioUrl]

diff --git a/docarray/typing/tensor/audio/audio_tensor.py b/docarray/typing/tensor/audio/audio_tensor.py
@@ -1,13 +1,23 @@
 from typing import Union
 
 from docarray.typing.tensor.audio.audio_ndarray import AudioNdArray
+from docarray.utils.misc import is_tf_available, is_torch_available
 
-try:
-    import torch  # noqa: F401
-except ImportError:
-    AudioTensor = AudioNdArray
-
-else:
+torch_available = is_torch_available()
+if torch_available:
     from docarray.typing.tensor.audio.audio_torch_tensor import AudioTorchTensor
 
+tf_available = is_tf_available()
+if tf_available:
+    from docarray.typing.tensor.audio.audio_tensorflow_tensor import (
+        AudioTensorFlowTensor as AudioTFTensor,
+    )
+
+
+AudioTensor = AudioNdArray
+if tf_available and torch_available:
+    AudioTensor = Union[AudioNdArray, AudioTorchTensor, AudioTFTensor]  # type: ignore
+elif tf_available:
+    AudioTensor = Union[AudioNdArray, AudioTFTensor]  # type: ignore
+elif torch_available:
     AudioTensor = Union[AudioNdArray, AudioTorchTensor]  # type: ignore