diff --git a/docs/snippets/all/descriptors/descr_custom_archetype.py b/docs/snippets/all/descriptors/descr_custom_archetype.py index c5c30eba9b32a..548cb4c0e8b19 100644 --- a/docs/snippets/all/descriptors/descr_custom_archetype.py +++ b/docs/snippets/all/descriptors/descr_custom_archetype.py @@ -5,50 +5,26 @@ from typing import Any, Iterable import numpy.typing as npt -import pyarrow as pa import rerun as rr # pip install rerun-sdk -class CustomPosition3DBatch(rr.ComponentBatchLike): - def __init__(self: Any, positions: npt.ArrayLike) -> None: - self.position = rr.components.Position3DBatch(positions) - - def component_descriptor(self) -> rr.ComponentDescriptor: - return rr.ComponentDescriptor( - "user.CustomPosition3D", - archetype_name="user.CustomPoints3D", - archetype_field_name="custom_positions", - ) - - def as_arrow_array(self) -> pa.Array: - return self.position.as_arrow_array() - - class CustomPoints3D(rr.AsComponents): def __init__(self: Any, positions: npt.ArrayLike, colors: npt.ArrayLike) -> None: - self.positions = CustomPosition3DBatch(positions) - self.colors = rr.components.ColorBatch(colors) + self.positions = rr.components.Position3DBatch(positions).with_descriptor( + rr.ComponentDescriptor( + "user.CustomPosition3D", + archetype_name="user.CustomPoints3D", + archetype_field_name="custom_positions", + ) + ) + self.colors = rr.components.ColorBatch(colors).or_with_descriptor_overrides( + archetype_name="user.CustomPoints3D", + archetype_field_name="colors", + ) def as_component_batches(self) -> Iterable[rr.ComponentBatchLike]: - return ( - [rr.IndicatorComponentBatch("user.CustomPoints3D")] - + [ - rr.DescribedComponentBatch( - self.positions, - self.positions.component_descriptor().or_with_overrides( - archetype_name="user.CustomPoints3D", archetype_field_name="custom_positions" - ), - ) - ] - + [ - rr.DescribedComponentBatch( - self.colors, - self.colors.component_descriptor().or_with_overrides( - archetype_name="user.CustomPoints3D", archetype_field_name="colors" - ), - ) - ] - ) + print([rr.IndicatorComponentBatch("user.CustomPoints3D"), self.positions, self.colors]) + return [rr.IndicatorComponentBatch("user.CustomPoints3D"), self.positions, self.colors] rr.init("rerun_example_descriptors_custom_archetype") diff --git a/docs/snippets/all/descriptors/descr_custom_component.py b/docs/snippets/all/descriptors/descr_custom_component.py index c119d0f587750..d4cd2d347641d 100644 --- a/docs/snippets/all/descriptors/descr_custom_component.py +++ b/docs/snippets/all/descriptors/descr_custom_component.py @@ -2,31 +2,18 @@ from __future__ import annotations -from typing import Any - -import numpy.typing as npt -import pyarrow as pa import rerun as rr # pip install rerun-sdk - -class CustomPosition3DBatch(rr.ComponentBatchLike): - def __init__(self: Any, positions: npt.ArrayLike) -> None: - self.position = rr.components.Position3DBatch(positions) - - def component_descriptor(self) -> rr.ComponentDescriptor: - return rr.ComponentDescriptor( - "user.CustomPosition3D", - archetype_name="user.CustomArchetype", - archetype_field_name="custom_positions", - ) - - def as_arrow_array(self) -> pa.Array: - return self.position.as_arrow_array() - - rr.init("rerun_example_descriptors_custom_component") rr.spawn() -rr.log("data", [CustomPosition3DBatch([[1, 2, 3]])], static=True) +positions = rr.components.Position3DBatch([1, 2, 3]).with_descriptor( + rr.ComponentDescriptor( + "user.CustomPosition3D", + archetype_name="user.CustomArchetype", + archetype_field_name="custom_positions", + ) +) +rr.log("data", [positions], static=True) # The tags are indirectly checked by the Rust version (have a look over there for more info). diff --git a/docs/snippets/all/tutorials/custom_data.py b/docs/snippets/all/tutorials/custom_data.py index 3eb172709ba12..907683eff8350 100755 --- a/docs/snippets/all/tutorials/custom_data.py +++ b/docs/snippets/all/tutorials/custom_data.py @@ -12,7 +12,7 @@ import rerun as rr -class ConfidenceBatch(rr.ComponentBatchLike): +class ConfidenceBatch(rr.ComponentBatchMixin): """A batch of confidence data.""" def __init__(self: Any, confidence: npt.ArrayLike) -> None: @@ -30,24 +30,17 @@ def as_arrow_array(self) -> pa.Array: class CustomPoints3D(rr.AsComponents): """A custom archetype that extends Rerun's builtin `Points3D` archetype with a custom component.""" - def __init__(self: Any, points3d: npt.ArrayLike, confidences: npt.ArrayLike) -> None: - self.points3d = points3d - self.confidences = confidences + def __init__(self: Any, positions: npt.ArrayLike, confidences: npt.ArrayLike) -> None: + self.points3d = rr.Points3D(positions) + self.confidences = ConfidenceBatch(confidences).or_with_descriptor_overrides( + archetype_name="user.CustomPoints3D", archetype_field_name="confidences" + ) def as_component_batches(self) -> Iterable[rr.ComponentBatchLike]: - points3d = np.asarray(self.points3d) - confidences = ConfidenceBatch(self.confidences) return ( - list(rr.Points3D(points3d).as_component_batches()) # The components from Points3D + list(self.points3d.as_component_batches()) # The components from Points3D + [rr.IndicatorComponentBatch("user.CustomPoints3D")] # Our custom indicator - + [ - rr.DescribedComponentBatch( - confidences, - confidences.component_descriptor().or_with_overrides( - archetype_name="user.CustomPoints3D", archetype_field_name="confidences" - ), - ) - ] # Custom confidence data + + [self.confidences] # Custom confidence data ) @@ -59,14 +52,14 @@ def log_custom_data() -> None: rr.log( "left/my_confident_point_cloud", CustomPoints3D( - points3d=point_grid, + positions=point_grid, confidences=[42], ), ) rr.log( "right/my_polarized_point_cloud", - CustomPoints3D(points3d=point_grid, confidences=np.arange(0, len(point_grid))), + CustomPoints3D(positions=point_grid, confidences=np.arange(0, len(point_grid))), ) diff --git a/rerun_py/rerun_sdk/rerun/__init__.py b/rerun_py/rerun_sdk/rerun/__init__.py index b3b973da50216..ffa7ac2821079 100644 --- a/rerun_py/rerun_sdk/rerun/__init__.py +++ b/rerun_py/rerun_sdk/rerun/__init__.py @@ -33,6 +33,7 @@ remote as remote, ) from ._baseclasses import ( + ComponentBatchMixin as ComponentBatchMixin, ComponentColumn as ComponentColumn, ComponentDescriptor as ComponentDescriptor, DescribedComponentBatch as DescribedComponentBatch, diff --git a/rerun_py/rerun_sdk/rerun/_baseclasses.py b/rerun_py/rerun_sdk/rerun/_baseclasses.py index e7bfc43a71d83..d1c43c6c9ec92 100644 --- a/rerun_py/rerun_sdk/rerun/_baseclasses.py +++ b/rerun_py/rerun_sdk/rerun/_baseclasses.py @@ -414,6 +414,44 @@ def component_descriptor(self) -> ComponentDescriptor: """ return self._COMPONENT_DESCRIPTOR # type: ignore[attr-defined, no-any-return] + def with_descriptor(self, descriptor: ComponentDescriptor) -> DescribedComponentBatch: + """Wraps the current `ComponentBatchLike` in a `DescribedComponentBatch` with the given descriptor.""" + return DescribedComponentBatch(self, descriptor) + + def with_descriptor_overrides( + self, *, archetype_name: str | None, archetype_field_name: str | None + ) -> ComponentBatchLike: + """Unconditionally sets `archetype_name` & `archetype_field_name` to the given ones (if specified).""" + descriptor = self.component_descriptor() + component_name = descriptor.component_name + archetype_name = archetype_name if archetype_name is not None else descriptor.archetype_name + archetype_field_name = ( + archetype_field_name if archetype_field_name is not None else descriptor.archetype_field_name + ) + return DescribedComponentBatch( + self, + ComponentDescriptor( + component_name, archetype_name=archetype_name, archetype_field_name=archetype_field_name + ), + ) + + def or_with_descriptor_overrides( + self, *, archetype_name: str | None, archetype_field_name: str | None + ) -> ComponentBatchLike: + """Sets `archetype_name` & `archetype_field_name` to the given one iff it's not already set.""" + descriptor = self.component_descriptor() + component_name = descriptor.component_name + archetype_name = descriptor.archetype_name if descriptor.archetype_name is not None else archetype_name + archetype_field_name = ( + descriptor.archetype_field_name if descriptor.archetype_field_name is not None else archetype_field_name + ) + return DescribedComponentBatch( + self, + ComponentDescriptor( + component_name, archetype_name=archetype_name, archetype_field_name=archetype_field_name + ), + ) + def partition(self, lengths: npt.ArrayLike) -> ComponentColumn: """ Partitions the component into multiple sub-batches. This wraps the inner arrow