diff --git a/rerun_py/rerun_sdk/rerun2/components/label_ext.py b/rerun_py/rerun_sdk/rerun2/components/label_ext.py new file mode 100644 index 0000000000000..11e970c47e433 --- /dev/null +++ b/rerun_py/rerun_sdk/rerun2/components/label_ext.py @@ -0,0 +1,16 @@ +from __future__ import annotations + +from typing import Any, Sequence + +import pyarrow as pa + + +class LabelArrayExt: + @staticmethod + def _from_similar(data: Any | None, *, mono: type, mono_aliases: type, many: type, many_aliases: type, arrow: type): + if isinstance(data, Sequence): + array = [str(datum) for datum in data] + else: + array = [str(data)] + + return arrow().wrap_array(pa.array(array, type=arrow().storage_type)) diff --git a/rerun_py/tests/unit/test_points2d.py b/rerun_py/tests/unit/test_points2d.py index f17af18d550e8..9beb0deb666a3 100644 --- a/rerun_py/tests/unit/test_points2d.py +++ b/rerun_py/tests/unit/test_points2d.py @@ -8,6 +8,8 @@ from rerun_sdk import rerun2 as rr from rerun_sdk.rerun2 import components as rrc +# TODO(cmc): roundtrips (serialize in python, deserialize in rust) + def test_points2d() -> None: points_arrays = [ @@ -49,27 +51,43 @@ def test_points2d() -> None: np.array([42, 43]), ] + # TODO: color + + labels_arrays = [ + None, + # LabelArrayLike: Sequence[LabelLike]: str + ["hello", "friend"], + # LabelArrayLike: Sequence[LabelLike]: Label + [ + rrc.Label("hello"), + rrc.Label("friend"), + ], + ] + all_permuted_arrays = list( itertools.product( # type: ignore[call-overload] *[ points_arrays, radii_arrays, + labels_arrays, ] ) ) - for points, radii in all_permuted_arrays: + for points, radii, labels in all_permuted_arrays: print( f"rr.Points2D(\n" f" {points}\n" f" radii={radii}\n" + f" labels={labels}\n" f")" ) - arch = rr.Points2D(points, radii=radii) + arch = rr.Points2D(points, radii=radii, labels=labels) + print(f"{arch}\n") assert arch.points == rrc.Point2DArray.from_similar([[1.0, 2.0], [3.0, 4.0]]) assert arch.radii == rrc.RadiusArray.from_similar([42, 43] if radii is not None else []) - print(f"{arch}\n") + assert arch.labels == rrc.LabelArray.from_similar(["hello", "friend"] if labels is not None else []) if __name__ == "__main__":