docarray · samsja · Mar 30, 2023 · Mar 23, 2023 · Mar 23, 2023 · Mar 24, 2023
diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml
@@ -112,6 +112,7 @@ jobs:
           python -m pip install --upgrade pip
           python -m pip install poetry
           poetry install --all-extras
+          sudo apt-get update
           sudo apt-get install --no-install-recommends ffmpeg
 
       - name: Test
@@ -157,6 +158,7 @@ jobs:
           python -m pip install poetry
           rm poetry.lock
           poetry install --all-extras
+          sudo apt-get update
           sudo apt-get install --no-install-recommends ffmpeg
 
       - name: Test
@@ -187,6 +189,7 @@ jobs:
           python -m pip install poetry
           poetry install --all-extras
           poetry run pip install protobuf==3.19.0 # we check that we support 3.19
+          sudo apt-get update
           sudo apt-get install --no-install-recommends ffmpeg
 
       - name: Test
@@ -214,6 +217,7 @@ jobs:
           python -m pip install --upgrade pip
           python -m pip install poetry
           poetry install --all-extras
+          sudo apt-get update
           sudo apt-get install --no-install-recommends ffmpeg
 
       - name: Test
@@ -242,6 +246,7 @@ jobs:
           poetry install --all-extras
           poetry run pip install protobuf==3.19.0
           poetry run pip install tensorflow==2.11.0
+          sudo apt-get update
           sudo apt-get install --no-install-recommends ffmpeg
 
       - name: Test

diff --git a/docarray/array/array/io.py b/docarray/array/array/io.py
@@ -32,6 +32,7 @@
     _dict_to_access_paths,
 )
 from docarray.utils._internal.compress import _decompress_bytes, _get_compress_ctx
+from docarray.utils._internal.misc import import_library
 
 if TYPE_CHECKING:
     import pandas as pd
@@ -500,7 +501,10 @@ def to_pandas(self) -> 'pd.DataFrame':
 
         :return: pandas.DataFrame
         """
-        import pandas as pd
+        if TYPE_CHECKING:
+            import pandas as pd
+        else:
+            pd = import_library('pandas', raise_error=True)
 
         fields = self.document_type._get_access_paths()
         df = pd.DataFrame(columns=fields)

diff --git a/docarray/array/array/sequence_indexing_mixin.py b/docarray/array/array/sequence_indexing_mixin.py
@@ -1,5 +1,6 @@
 import abc
 from typing import (
+    TYPE_CHECKING,
     Any,
     Iterable,
     MutableSequence,
@@ -14,7 +15,7 @@
 
 import numpy as np
 
-from docarray.utils._internal.misc import is_torch_available
+from docarray.utils._internal.misc import import_library
 
 T_item = TypeVar('T_item')
 T = TypeVar('T', bound='IndexingSequenceMixin')
@@ -85,11 +86,11 @@ def _normalize_index_item(
             return item.tolist()
 
         # torch index types
-        torch_available = is_torch_available()
-        if torch_available:
+        if TYPE_CHECKING:
             import torch
         else:
-            raise ValueError(f'Invalid index type {type(item)}')
+            torch = import_library('torch', raise_error=True)
+
         allowed_torch_dtypes = [
             torch.bool,
             torch.int64,

diff --git a/docarray/base_doc/__init__.py b/docarray/base_doc/__init__.py
@@ -1,6 +1,24 @@
 from docarray.base_doc.any_doc import AnyDoc
 from docarray.base_doc.base_node import BaseNode
 from docarray.base_doc.doc import BaseDoc
-from docarray.base_doc.doc_response import DocResponse
+from docarray.utils._internal.misc import (
+    _get_path_from_docarray_root_level,
+    import_library,
+)
 
-__all__ = ['AnyDoc', 'BaseDoc', 'BaseNode', 'DocResponse']
+__all__ = ['AnyDoc', 'BaseDoc', 'BaseNode']
+
+
+def __getattr__(name: str):
+    if name == 'DocResponse':
+        import_library('fastapi', raise_error=True)
+        from docarray.base_doc.doc_response import DocResponse
+
+        if name not in __all__:
+            __all__.append(name)
+
+        return DocResponse
+    else:
+        raise ImportError(
+            f'cannot import name \'{name}\' from \'{_get_path_from_docarray_root_level(__file__)}\''
+        )
diff --git a/docarray/base_doc/doc_response.py b/docarray/base_doc/doc_response.py
@@ -1,14 +1,12 @@
-from typing import Any
+from typing import TYPE_CHECKING, Any
 
-try:
-    from fastapi.responses import JSONResponse, Response
-except ImportError:
+from docarray.utils._internal.misc import import_library
 
-    class NoImportResponse:
-        def __init__(self, *args, **kwargs):
-            ImportError('fastapi is not installed')
-
-    Response = JSONResponse = NoImportResponse  # type: ignore
+if TYPE_CHECKING:
+    from fastapi.responses import JSONResponse
+else:
+    fastapi = import_library('fastapi', raise_error=True)
+    JSONResponse = fastapi.responses.JSONResponse
 
 
 class DocResponse(JSONResponse):

diff --git a/docarray/base_doc/mixins/io.py b/docarray/base_doc/mixins/io.py
@@ -21,24 +21,23 @@
 from docarray.typing import NdArray
 from docarray.typing.proto_register import _PROTO_TYPE_NAME_TO_CLASS
 from docarray.utils._internal.compress import _compress_bytes, _decompress_bytes
-from docarray.utils._internal.misc import is_tf_available, is_torch_available
+from docarray.utils._internal.misc import import_library
 
-tf_available = is_tf_available()
-if tf_available:
+if TYPE_CHECKING:
     import tensorflow as tf  # type: ignore
-
-    from docarray.typing import TensorFlowTensor
-
-torch_available = is_torch_available()
-if torch_available:
     import torch
-
-    from docarray.typing import TorchTensor
-
-if TYPE_CHECKING:
     from pydantic.fields import ModelField
 
     from docarray.proto import DocumentProto, NodeProto
+    from docarray.typing import TensorFlowTensor, TorchTensor
+else:
+    tf = import_library('tensorflow', raise_error=False)
+    if tf is not None:
+        from docarray.typing import TensorFlowTensor
+
+    torch = import_library('torch', raise_error=False)
+    if torch is not None:
+        from docarray.typing import TorchTensor
 
 
 T = TypeVar('T', bound='IOMixin')
@@ -68,12 +67,12 @@ def _type_to_protobuf(value: Any) -> 'NodeProto':
         return nested_item
 
     base_node_wrap: BaseNode
-    if torch_available:
+    if torch is not None:
         if isinstance(value, torch.Tensor):
             base_node_wrap = TorchTensor._docarray_from_native(value)
             return base_node_wrap._to_node_protobuf()
 
-    if tf_available:
+    if tf is not None:
         if isinstance(value, tf.Tensor):
             base_node_wrap = TensorFlowTensor._docarray_from_native(value)
             return base_node_wrap._to_node_protobuf()

diff --git a/docarray/computation/tensorflow_backend.py b/docarray/computation/tensorflow_backend.py
@@ -1,13 +1,19 @@
 import typing
-from typing import Callable, List, Optional, Tuple
+from typing import TYPE_CHECKING, Callable, List, Optional, Tuple
 
 import numpy as np
-import tensorflow as tf  # type: ignore
-import tensorflow._api.v2.experimental.numpy as tnp  # type: ignore
 
 from docarray.computation import AbstractComputationalBackend
 from docarray.computation.abstract_numpy_based_backend import AbstractNumpyBasedBackend
 from docarray.typing import TensorFlowTensor
+from docarray.utils._internal.misc import import_library
+
+if TYPE_CHECKING:
+    import tensorflow as tf  # type: ignore
+    import tensorflow._api.v2.experimental.numpy as tnp  # type: ignore
+else:
+    tf = import_library('tensorflow', raise_error=True)
+    tnp = tf._api.v2.experimental.numpy
 
 
 def _unsqueeze_if_single_axis(*matrices: tf.Tensor) -> List[tf.Tensor]:

diff --git a/docarray/computation/torch_backend.py b/docarray/computation/torch_backend.py
@@ -1,9 +1,14 @@
-from typing import Any, List, Optional, Tuple, Union
+from typing import TYPE_CHECKING, Any, List, Optional, Tuple, Union
 
 import numpy as np
-import torch
 
 from docarray.computation.abstract_comp_backend import AbstractComputationalBackend
+from docarray.utils._internal.misc import import_library
+
+if TYPE_CHECKING:
+    import torch
+else:
+    torch = import_library('torch', raise_error=True)
 
 
 def _unsqueeze_if_single_axis(*matrices: torch.Tensor) -> List[torch.Tensor]:

diff --git a/docarray/documents/audio.py b/docarray/documents/audio.py
@@ -1,4 +1,4 @@
-from typing import Any, Optional, Type, TypeVar, Union
+from typing import TYPE_CHECKING, Any, Optional, Type, TypeVar, Union
 
 import numpy as np
 
@@ -7,15 +7,14 @@
 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._internal.misc import is_tf_available, is_torch_available
+from docarray.utils._internal.misc import import_library
 
-torch_available = is_torch_available()
-if torch_available:
-    import torch
-
-tf_available = is_tf_available()
-if tf_available:
+if TYPE_CHECKING:
     import tensorflow as tf  # type: ignore
+    import torch
+else:
+    torch = import_library('torch', raise_error=False)
+    tf = import_library('tensorflow', raise_error=False)
 
 
 T = TypeVar('T', bound='AudioDoc')
@@ -110,9 +109,9 @@ def validate(
         if isinstance(value, str):
             value = cls(url=value)
         elif isinstance(value, (AbstractTensor, np.ndarray)) or (
-            torch_available
+            torch is not None
             and isinstance(value, torch.Tensor)
-            or (tf_available and isinstance(value, tf.Tensor))
+            or (tf is not None and isinstance(value, tf.Tensor))
         ):
             value = cls(tensor=value)
 

diff --git a/docarray/documents/image.py b/docarray/documents/image.py
@@ -1,22 +1,21 @@
-from typing import Any, Optional, Type, TypeVar, Union
+from typing import TYPE_CHECKING, Any, Optional, Type, TypeVar, Union
 
 import numpy as np
 
 from docarray.base_doc import BaseDoc
 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._internal.misc import is_tf_available, is_torch_available
+from docarray.utils._internal.misc import import_library
 
-T = TypeVar('T', bound='ImageDoc')
-
-torch_available = is_torch_available()
-if torch_available:
+if TYPE_CHECKING:
+    import tensorflow as tf  # type: ignore
     import torch
+else:
+    tf = import_library('tensorflow', raise_error=False)
+    torch = import_library('torch', raise_error=False)
 
-tf_available = is_tf_available()
-if tf_available:
-    import tensorflow as tf  # type: ignore
+T = TypeVar('T', bound='ImageDoc')
 
 
 class ImageDoc(BaseDoc):
@@ -99,8 +98,8 @@ def validate(
             value = cls(url=value)
         elif (
             isinstance(value, (AbstractTensor, np.ndarray))
-            or (torch_available and isinstance(value, torch.Tensor))
-            or (tf_available and isinstance(value, tf.Tensor))
+            or (torch is not None and isinstance(value, torch.Tensor))
+            or (tf is not None and isinstance(value, tf.Tensor))
         ):
             value = cls(tensor=value)
         elif isinstance(value, bytes):

diff --git a/docarray/documents/mesh/vertices_and_faces.py b/docarray/documents/mesh/vertices_and_faces.py
@@ -1,7 +1,8 @@
-from typing import Any, Type, TypeVar, Union
+from typing import TYPE_CHECKING, Any, Type, TypeVar, Union
 
 from docarray.base_doc import BaseDoc
 from docarray.typing.tensor.tensor import AnyTensor
+from docarray.utils._internal.misc import import_library
 
 T = TypeVar('T', bound='VerticesAndFaces')
 
@@ -28,9 +29,12 @@ def validate(
     def display(self) -> None:
         """
         Plot mesh consisting of vertices and faces.
-        To use this you need to install trimesh[easy]: `pip install 'trimesh[easy]'`.
         """
-        import trimesh
+        if TYPE_CHECKING:
+            import trimesh
+        else:
+            trimesh = import_library('trimesh', raise_error=True)
+
         from IPython.display import display
 
         if self.vertices is None or self.faces is None:

diff --git a/docarray/documents/point_cloud/point_cloud_3d.py b/docarray/documents/point_cloud/point_cloud_3d.py
@@ -1,20 +1,20 @@
-from typing import Any, Optional, Type, TypeVar, Union
+from typing import TYPE_CHECKING, Any, Optional, Type, TypeVar, Union
 
 import numpy as np
 
 from docarray.base_doc import BaseDoc
 from docarray.documents.point_cloud.points_and_colors import PointsAndColors
 from docarray.typing import AnyEmbedding, PointCloud3DUrl
 from docarray.typing.tensor.abstract_tensor import AbstractTensor
-from docarray.utils._internal.misc import is_tf_available, is_torch_available
+from docarray.utils._internal.misc import import_library
 
-torch_available = is_torch_available()
-if torch_available:
+if TYPE_CHECKING:
+    import tensorflow as tf  # type: ignore
     import torch
+else:
+    tf = import_library('tensorflow', raise_error=False)
+    torch = import_library('torch', raise_error=False)
 
-tf_available = is_tf_available()
-if tf_available:
-    import tensorflow as tf  # type: ignore
 
 T = TypeVar('T', bound='PointCloud3D')
 
@@ -124,9 +124,9 @@ def validate(
         if isinstance(value, str):
             value = cls(url=value)
         elif isinstance(value, (AbstractTensor, np.ndarray)) or (
-            torch_available
+            torch is not None
             and isinstance(value, torch.Tensor)
-            or (tf_available and isinstance(value, tf.Tensor))
+            or (tf is not None and isinstance(value, tf.Tensor))
         ):
             value = cls(tensors=PointsAndColors(points=value))