update python tests/extensions accordingly

rerun-io · Jun 27, 2023 · 9d91529 · 9d91529
1 parent 29277fd
commit 9d91529
Show file tree

Hide file tree

Showing 2 changed files with 10 additions and 10 deletions.
diff --git a/rerun_py/rerun_sdk/rerun/_rerun2/components/point2d_ext.py b/rerun_py/rerun_sdk/rerun/_rerun2/components/point2d_ext.py
@@ -14,10 +14,15 @@ def _from_similar(
         data: Any | None, *, mono: type, mono_aliases: Any, many: type, many_aliases: Any, arrow: type
     ) -> pa.Array:
         if isinstance(data, Sequence) and (len(data) > 0 and isinstance(data[0], mono)):
-            points = np.concatenate([np.asarray(datum.position, dtype=np.float32) for datum in data])
+            points = np.concatenate([np.asarray([datum.x, datum.y], dtype=np.float32) for datum in data])
         else:
             points = np.asarray(data, dtype=np.float32)
 
-        points = points.reshape((-1,))
+        points = points.reshape((-1, 2))
 
-        return arrow().wrap_array(pa.FixedSizeListArray.from_arrays(points, type=arrow().storage_type))
+        return arrow().wrap_array(
+            pa.StructArray.from_arrays(
+                arrays=[pa.array(c, type=pa.float32()) for c in points.T],
+                fields=list(arrow().storage_type),
+            )
+        )
diff --git a/rerun_py/tests/unit/test_points2d.py b/rerun_py/tests/unit/test_points2d.py
@@ -22,8 +22,8 @@ def test_points2d() -> None:
         np.array([]),
         # Point2DArrayLike: Sequence[Point2DLike]: Point2D
         [
-            rrc.Point2D([1, 2]),
-            rrc.Point2D([3, 4]),
+            rrc.Point2D(1, 2),
+            rrc.Point2D(3, 4),
         ],
         # Point2DArrayLike: Sequence[Point2DLike]: npt.NDArray[np.float32]
         [
@@ -36,11 +36,6 @@ def test_points2d() -> None:
         [1, 2, 3, 4],
         # Point2DArrayLike: npt.NDArray[np.float32]
         np.array([[1, 2], [3, 4]], dtype=np.float32),
-        # Point2DArrayLike: Sequence[Point2DLike]
-        [
-            rrc.Point2D([1, 2]),
-            rrc.Point2D([3, 4]),
-        ],
         # Point2DArrayLike: npt.NDArray[np.float32]
         np.array([1, 2, 3, 4], dtype=np.float32),
     ]