python helpers

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")

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).

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))),
     )
 
 

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,

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