fix: mypy and tests

Signed-off-by: anna-charlotte <[email protected]>
docarray · samsja · Feb 17, 2023 · Feb 9, 2023 · Feb 9, 2023 · Feb 9, 2023
commit db42712db07cfedb0cca268da7b7c3d962c517ea
diff --git a/docarray/documents/mesh/__init__.py b/docarray/documents/mesh/__init__.py
@@ -1,3 +1,3 @@
-from docarray.documents.mesh.mesh import Mesh3D
+from docarray.documents.mesh.mesh_3d import Mesh3D
 
 __all__ = ['Mesh3D']
diff --git a/docarray/documents/mesh/mesh.py → docarray/documents/mesh/mesh_3d.py b/docarray/documents/mesh/mesh.py → docarray/documents/mesh/mesh_3d.py
diff --git a/docarray/documents/mesh/vertices_and_faces.py b/docarray/documents/mesh/vertices_and_faces.py
@@ -7,7 +7,13 @@
 
 
 class VerticesAndFaces(BaseDocument):
-    """ """
+    """
+    Document for handling 3D mesh tensor data.
+
+    A VerticesAndFaces Document can contain an AnyTensor containing the vertices
+    information (`VerticesAndFaces.vertices`), and an AnyTensor containing the faces
+    information (`VerticesAndFaces.faces`).
+    """
 
     vertices: Optional[AnyTensor]
     faces: Optional[AnyTensor]

diff --git a/docarray/documents/point_cloud/__init__.py b/docarray/documents/point_cloud/__init__.py
@@ -1,3 +1,3 @@
-from docarray.documents.point_cloud.point_cloud import PointCloud3D
+from docarray.documents.point_cloud.point_cloud_3d import PointCloud3D
 
 __all__ = ['PointCloud3D']
diff --git a/...rray/documents/point_cloud/point_cloud.py → ...y/documents/point_cloud/point_cloud_3d.py b/...rray/documents/point_cloud/point_cloud.py → ...y/documents/point_cloud/point_cloud_3d.py
@@ -28,8 +28,9 @@ class PointCloud3D(BaseDocument):
     representation, the point cloud is a fixed size ndarray (shape=(n_samples, 3)) and
     hence easier for deep learning algorithms to handle.
 
-    A PointCloud3D Document can contain an PointCloud3DUrl (`PointCloud3D.url`), an
-    AnyTensor (`PointCloud3D.tensor`), and an AnyEmbedding (`PointCloud3D.embedding`).
+    A PointCloud3D Document can contain an PointCloud3DUrl (`PointCloud3D.url`),
+    a PointsAndColors object (`PointCloud3D.tensors`), and an AnyEmbedding
+    (`PointCloud3D.embedding`).
 
     EXAMPLE USAGE:
 

diff --git a/docarray/documents/point_cloud/points_and_colors.py b/docarray/documents/point_cloud/points_and_colors.py
@@ -20,7 +20,13 @@
 
 
 class PointsAndColors(BaseDocument):
-    """ """
+    """
+    Document for handling point clouds tensor data.
+
+    A PointsAndColors Document can contain an AnyTensor containing the points in
+    3D space information (`PointsAndColors.points`), and an AnyTensor containing
+    the points' color information (`PointsAndColors.colors`).
+    """
 
     points: AnyTensor
     colors: Optional[AnyTensor]

diff --git a/tests/units/typing/url/test_point_cloud_url.py b/tests/units/typing/url/test_point_cloud_url.py
@@ -49,11 +49,11 @@ def test_load(file_format, file_path):
 def test_load_with_multiple_geometries_true(file_format, file_path):
     n_samples = 100
     url = parse_obj_as(PointCloud3DUrl, file_path)
-    point_cloud = url.load(samples=n_samples, multiple_geometries=True)
+    tensors = url.load(samples=n_samples, multiple_geometries=True)
 
-    assert isinstance(point_cloud, np.ndarray)
-    assert len(point_cloud.shape) == 3
-    assert point_cloud.shape[1:] == (100, 3)
+    assert isinstance(tensors.points, np.ndarray)
+    assert len(tensors.points.shape) == 3
+    assert tensors.points.shape[1:] == (100, 3)
 
 
 def test_json_schema():