Files
rerun-io--rerun/rerun_py/tests/unit/test_tensor_slice_selection.py
2026-07-13 13:05:14 +08:00

89 lines
2.9 KiB
Python

from __future__ import annotations
import itertools
from typing import TYPE_CHECKING, cast
import numpy as np
import rerun as rr
import rerun.blueprint as rrb
from .common_arrays import none_empty_or_value
if TYPE_CHECKING:
from rerun.blueprint.datatypes import TensorDimensionIndexSliderArrayLike
def test_tensor_slice_selection() -> None:
widths = [
None,
2,
rr.datatypes.TensorDimensionSelection(dimension=2, invert=False),
rr.components.TensorWidthDimension(dimension=2, invert=False),
]
heights = [
None,
3,
rr.datatypes.TensorDimensionSelection(dimension=3, invert=False),
rr.components.TensorHeightDimension(dimension=3, invert=False),
]
indices_arrays = [
[
rr.components.TensorDimensionIndexSelection(dimension=1, index=3),
rr.components.TensorDimensionIndexSelection(dimension=2, index=2),
rr.components.TensorDimensionIndexSelection(dimension=3, index=1),
],
None,
]
slider_arrays = [
None,
[1, 2, 3],
[
rrb.components.TensorDimensionIndexSlider(1),
rrb.components.TensorDimensionIndexSlider(2),
rrb.components.TensorDimensionIndexSlider(3),
],
np.array([1, 2, 3]),
]
all_arrays = itertools.zip_longest(
widths,
heights,
indices_arrays,
slider_arrays,
)
for width, height, indices, slider in all_arrays:
width = cast("rr.datatypes.TensorDimensionSelectionLike | None", width)
height = cast("rr.datatypes.TensorDimensionSelectionLike | None", height)
indices = cast("rr.datatypes.TensorDimensionIndexSelectionArrayLike | None", indices)
slider = cast("TensorDimensionIndexSliderArrayLike | None", slider)
print(
f"rr.TensorSliceSelection(\n"
f" width={width!r}\n"
f" height={height!r}\n"
f" indices={indices!r}\n"
f" slider={slider!r}\n"
f")",
)
arch = rrb.TensorSliceSelection(
width=width,
height=height,
indices=indices,
slider=slider,
)
print(f"{arch}\n")
assert arch.width == rr.components.TensorWidthDimensionBatch._converter(
none_empty_or_value(width, rr.components.TensorWidthDimension(dimension=2, invert=False)),
)
assert arch.height == rr.components.TensorHeightDimensionBatch._converter(
none_empty_or_value(height, rr.components.TensorHeightDimension(dimension=3, invert=False)),
)
assert arch.indices == rr.components.TensorDimensionIndexSelectionBatch._converter(
none_empty_or_value(indices, indices_arrays[0]),
)
assert arch.slider == rrb.components.TensorDimensionIndexSliderBatch._converter(
none_empty_or_value(slider, [1, 2, 3]),
)