fix: mypy

Signed-off-by: anna-charlotte <[email protected]>
docarray · samsja · Feb 17, 2023 · Feb 9, 2023 · Feb 9, 2023 · Feb 9, 2023
commit 75b33d8e340aa6f98748cb4e7f01263ffdd7a9fb
diff --git a/docarray/documents/point_cloud.py b/docarray/documents/point_cloud.py
@@ -127,27 +127,24 @@ def display(self, display_from: str = 'url', samples: int = 10000) -> None:
         if display_from not in ['tensor', 'url']:
             raise ValueError(f'Expected one of ["tensor", "url"], got "{display_from}"')
 
-        if getattr(self, display_from) is None:
-            raise ValueError(
-                f'Can not to display point cloud from {display_from} when the '
-                f'{display_from} is None.'
-            )
-
         if display_from == 'url':
+            if self.url is None:
+                raise ValueError(
+                    'Can\'t display point cloud from url when url is None.'
+                )
             tensor = self.url.load(samples=samples)
             colors = np.tile(
                 np.array([0, 0, 0]), (tensor.get_comp_backend().shape(tensor)[0], 1)
             )
         else:
+            if self.tensor is None:
+                raise ValueError('Can\'t display mesh from tensor when tensor is None.')
             tensor = self.tensor
             comp_be = self.tensor.get_comp_backend()
             colors = (
                 self.color_tensor
                 if self.color_tensor
-                else np.tile(
-                    np.array([0, 0, 0]),
-                    (comp_be.shape(tensor)[0], 1),
-                )
+                else np.tile(np.array([0, 0, 0]), (comp_be.shape(tensor)[0], 1))
             )
 
         pc = trimesh.points.PointCloud(vertices=tensor, colors=colors)

diff --git a/docarray/typing/tensor/tensor.py b/docarray/typing/tensor/tensor.py
@@ -13,11 +13,10 @@
     from docarray.typing.tensor.tensorflow_tensor import TensorFlowTensor  # noqa: F401
 
 
+AnyTensor = Union[NdArray]
 if torch_available and tf_available:
-    AnyTensor = Union[NdArray, TorchTensor, TensorFlowTensor]
+    AnyTensor = Union[NdArray, TorchTensor, TensorFlowTensor]  # type: ignore
 elif torch_available:
     AnyTensor = Union[NdArray, TorchTensor]  # type: ignore
 elif tf_available:
     AnyTensor = Union[NdArray, TensorFlowTensor]  # type: ignore
-else:
-    AnyTensor = Union[NdArray]  # type: ignore