Add python unit tests for GeoPoints

rerun_py/tests/unit/common_arrays.py

@@ -19,6 +19,9 @@
 )
 from rerun.datatypes import (
     Angle,
+    DVec2D,
+    DVec2DArrayLike,
+    DVec2DBatch,
     Float32ArrayLike,
     Quaternion,
     QuaternionArrayLike,
@@ -63,6 +66,45 @@ def none_empty_or_value(obj: Any, value: Any) -> Any:
         return value
 
 
+dvec2ds_arrays: list[DVec2DArrayLike] = [
+    [],
+    np.array([]),
+    # Vec2DArrayLike: Sequence[Point2DLike]:
+    [
+        DVec2D([1, 2]),
+        DVec2D([3, 4]),
+    ],
+    # Vec2DArrayLike: Sequence[Point2DLike]: npt.NDArray[np.float64]
+    [
+        np.array([1, 2], dtype=np.float64),
+        np.array([3, 4], dtype=np.float64),
+    ],
+    # Vec2DArrayLike: Sequence[Point2DLike]: Tuple[float, float]
+    [(1, 2), (3, 4)],
+    # Vec2DArrayLike: torch.tensor is np.ArrayLike
+    torch.tensor([(1, 2), (3, 4)], dtype=torch.float64),
+    # Vec2DArrayLike: Sequence[Point2DLike]: Sequence[float]
+    [1, 2, 3, 4],
+    # Vec2DArrayLike: npt.NDArray[np.float64]
+    np.array([[1, 2], [3, 4]], dtype=np.float64),
+    # Vec2DArrayLike: npt.NDArray[np.float64]
+    np.array([1, 2, 3, 4], dtype=np.float64),
+    # Vec2DArrayLike: npt.NDArray[np.float64]
+    np.array([1, 2, 3, 4], dtype=np.float64).reshape((2, 2, 1, 1, 1)),
+    # PyTorch array
+    torch.asarray([1, 2, 3, 4], dtype=torch.float64),
+]
+
+
+def dvec2ds_expected(obj: Any, type_: Any | None = None) -> Any:
+    if type_ is None:
+        type_ = DVec2DBatch
+
+    expected = none_empty_or_value(obj, [[1.0, 2.0], [3.0, 4.0]])
+
+    return type_._optional(expected)
+
+
 vec2ds_arrays: list[Vec2DArrayLike] = [
     [],
     np.array([]),

rerun_py/tests/unit/test_geopoints.py

@@ -0,0 +1,104 @@
+from __future__ import annotations
+
+import itertools
+from typing import Optional, cast
+
+import numpy as np
+import pytest
+import rerun as rr
+from rerun.components import (
+    Color,
+    ColorBatch,
+    LatLonBatch,
+)
+from rerun.datatypes import (
+    Rgba32ArrayLike,
+    Float32ArrayLike,
+    DVec2DArrayLike,
+)
+
+from .common_arrays import (
+    colors_arrays,
+    colors_expected,
+    dvec2ds_arrays as positions_arrays,
+    dvec2ds_expected as positions_expected,
+    radii_arrays,
+    radii_expected,
+)
+
+
+def test_geopoints() -> None:
+    all_arrays = itertools.zip_longest(
+        positions_arrays,
+        radii_arrays,
+        colors_arrays,
+    )
+
+    for positions, radii, colors in all_arrays:
+        positions = positions if positions is not None else positions_arrays[-1]
+
+        # make Pyright happy as it's apparently not able to track typing info through zip_longest
+        positions = cast(DVec2DArrayLike, positions)
+        radii = cast(Optional[Float32ArrayLike], radii)
+        colors = cast(Optional[Rgba32ArrayLike], colors)
+
+        print(f"rr.GeoPoints(\n    {positions}\n    radii={radii!r}\n    colors={colors!r}\n)")
+        arch = rr.GeoPoints(
+            positions,
+            radii=radii,
+            colors=colors,
+        )
+        print(f"{arch}\n")
+
+        assert arch.positions == positions_expected(positions, LatLonBatch)
+        assert arch.radii == radii_expected(radii)
+        assert arch.colors == colors_expected(colors)
+
+
+@pytest.mark.parametrize(
+    "data",
+    [
+        [0, 128, 0, 255],
+        [0, 128, 0],
+        np.array((0, 128, 0, 255)),
+        [0.0, 0.5, 0.0, 1.0],
+        np.array((0.0, 0.5, 0.0, 1.0)),
+    ],
+)
+def test_geopoint_single_color(data: Rgba32ArrayLike) -> None:
+    pts = rr.GeoPoints(np.zeros((5, 2)), colors=data)
+
+    assert pts.colors == ColorBatch(Color([0, 128, 0, 255]))
+
+
+@pytest.mark.parametrize(
+    "data",
+    [
+        [[0, 128, 0, 255], [128, 0, 0, 255]],
+        [[0, 128, 0], [128, 0, 0]],
+        np.array([[0, 128, 0, 255], [128, 0, 0, 255]]),
+        np.array([0, 128, 0, 255, 128, 0, 0, 255], dtype=np.uint8),
+        np.array([8388863, 2147483903], dtype=np.uint32),
+        np.array([[0, 128, 0], [128, 0, 0]]),
+        [[0.0, 0.5, 0.0, 1.0], [0.5, 0.0, 0.0, 1.0]],
+        [[0.0, 0.5, 0.0], [0.5, 0.0, 0.0]],
+        np.array([[0.0, 0.5, 0.0, 1.0], [0.5, 0.0, 0.0, 1.0]]),
+        np.array([[0.0, 0.5, 0.0], [0.5, 0.0, 0.0]]),
+        np.array([0.0, 0.5, 0.0, 1.0, 0.5, 0.0, 0.0, 1.0]),
+        # Note: Sequence[int] is interpreted as a single color when they are 3 or 4 long. For other lengths, they
+        # are interpreted as list of packed uint32 colors. Note that this means one cannot pass an len=N*4 flat list of
+        # color components.
+        [8388863, 2147483903],
+    ],
+)
+def test_point2d_multiple_colors(data: Rgba32ArrayLike) -> None:
+    pts = rr.GeoPoints(np.zeros((5, 2)), colors=data)
+
+    assert pts.colors == ColorBatch([
+        Color([0, 128, 0, 255]),
+        Color([128, 0, 0, 255]),
+    ])
+
+
+if __name__ == "__main__":
+    test_geopoints()

rerun_py/tests/unit/test_points2d.py

@@ -52,7 +52,7 @@ def test_points2d() -> None:
     for positions, radii, colors, labels, draw_order, class_ids, keypoint_ids in all_arrays:
         positions = positions if positions is not None else positions_arrays[-1]
 
-        # make Pyright happy as it's apparently not able to track typing info trough zip_longest
+        # make Pyright happy as it's apparently not able to track typing info through zip_longest
         positions = cast(Vec2DArrayLike, positions)
         radii = cast(Optional[Float32ArrayLike], radii)
         colors = cast(Optional[Rgba32ArrayLike], colors)