fix: apply samis suggestion

anna-charlotte · anna-charlotte · commit e0e8d1a292fc · 2023-02-10T16:32:42.000+01:00
Signed-off-by: anna-charlotte &lt;charlotte.gerhaher@jina.ai&gt;
diff --git a/docarray/documents/point_cloud/points_and_colors.py b/docarray/documents/point_cloud/points_and_colors.py
@@ -3,10 +3,9 @@
 import numpy as np
 
 from docarray.base_document import BaseDocument
+from docarray.typing import AnyTensor
 from docarray.typing.tensor.abstract_tensor import AbstractTensor
-from docarray.typing.tensor.tensor import AnyTensor
-from docarray.typing.url.url_3d.point_cloud_url import _display_point_cloud
-from docarray.utils.misc import is_tf_available, is_torch_available
+from docarray.utils.misc import is_notebook, is_tf_available, is_torch_available
 
 torch_available = is_torch_available()
 if torch_available:
@@ -49,8 +48,21 @@ def display(self) -> None:
         """
         Plot point cloud consisting of points in 3D space and optionally colors.
         """
-        if self.points is None:
-            raise ValueError(
-                'Can\'t display point cloud from tensors when the points are None.'
+        import trimesh
+        from IPython.display import display
+
+        colors = (
+            self.colors
+            if self.colors is not None
+            else np.tile(
+                np.array([0, 0, 0]),
+                (self.points.get_comp_backend().shape(self.points)[0], 1),
             )
-        _display_point_cloud(points=self.points, colors=self.colors)
+        )
+        pc = trimesh.points.PointCloud(vertices=self.points, colors=colors)
+
+        if is_notebook():
+            s = trimesh.Scene(geometry=pc)
+            display(s.show())
+        else:
+            display(pc.show())
diff --git a/docarray/typing/url/url_3d/point_cloud_url.py b/docarray/typing/url/url_3d/point_cloud_url.py
@@ -1,18 +1,16 @@
-from typing import TYPE_CHECKING, Optional, TypeVar
+from typing import TYPE_CHECKING, TypeVar
 
 import numpy as np
 from pydantic import parse_obj_as
 
-from docarray.typing import AnyTensor
 from docarray.typing.proto_register import _register_proto
-from docarray.typing.tensor.abstract_tensor import AbstractTensor
 from docarray.typing.tensor.ndarray import NdArray
 from docarray.typing.url.url_3d.url_3d import Url3D
-from docarray.utils.misc import is_notebook
 
 if TYPE_CHECKING:
     from docarray.documents.point_cloud.points_and_colors import PointsAndColors
 
+
 T = TypeVar('T', bound='PointCloud3DUrl')
 
 
@@ -78,30 +76,4 @@ def display(self, samples: int = 10000) -> None:
         Plot point cloud from url.
         :param samples: number of points to sample from the mesh.
         """
-        tensors = self.load(samples=samples)
-        _display_point_cloud(points=tensors.points)
-
-
-def _display_point_cloud(
-    points: AnyTensor, colors: Optional[AbstractTensor] = None
-) -> None:
-    """
-    Plot point cloud from tensors.
-    :param points: tensor representing the point in 3D space, shape (n_points, 3).
-    :param colors: tensor representing the colors as RGB or RGB-A values,
-        shape (n_points, 3) or (n_points, 4).
-    """
-    import trimesh
-    from IPython.display import display
-
-    if colors is None:
-        colors = np.tile(
-            np.array([0, 0, 0]), (points.get_comp_backend().shape(points)[0], 1)
-        )
-    pc = trimesh.points.PointCloud(vertices=points, colors=colors)
-
-    if is_notebook():
-        s = trimesh.Scene(geometry=pc)
-        display(s.show())
-    else:
-        display(pc.show())
+        self.load(samples=samples).display()
