diff --git a/crates/re_types/source_hash.txt b/crates/re_types/source_hash.txt index 6224067380c5..2edfdc1adebe 100644 --- a/crates/re_types/source_hash.txt +++ b/crates/re_types/source_hash.txt @@ -1,4 +1,4 @@ # This is a sha256 hash for all direct and indirect dependencies of this crate's build script. # It can be safely removed at anytime to force the build script to run again. # Check out build.rs to see how it's computed. -afc01539cb778ef699e8cef27436420fe8dc4ee078fdd56d874e2409b3749d88 \ No newline at end of file +2482cd3e136880416056a88c296bd419d5b01626d2db3914e5e375c614123bcc \ No newline at end of file diff --git a/crates/re_types_builder/src/codegen/python.rs b/crates/re_types_builder/src/codegen/python.rs index 144570f4d58a..4cca28137f1f 100644 --- a/crates/re_types_builder/src/codegen/python.rs +++ b/crates/re_types_builder/src/codegen/python.rs @@ -41,6 +41,7 @@ impl CodeGenerator for PythonCodeGenerator { datatypes_path, arrow_registry, objs, + ObjectKind::Datatype, &objs.ordered_objects(ObjectKind::Datatype.into()), ) .0, @@ -55,6 +56,7 @@ impl CodeGenerator for PythonCodeGenerator { components_path, arrow_registry, objs, + ObjectKind::Component, &objs.ordered_objects(ObjectKind::Component.into()), ) .0, @@ -68,6 +70,7 @@ impl CodeGenerator for PythonCodeGenerator { archetypes_path, arrow_registry, objs, + ObjectKind::Archetype, &objs.ordered_objects(ObjectKind::Archetype.into()), ); filepaths.extend(paths); @@ -117,6 +120,7 @@ fn quote_objects( out_path: impl AsRef, arrow_registry: &ArrowRegistry, all_objects: &Objects, + kind: ObjectKind, objs: &[&Object], ) -> (Vec, Vec) { let out_path = out_path.as_ref(); @@ -176,6 +180,11 @@ fn quote_objects( code.push_text(&format!("# {AUTOGEN_WARNING}"), 2, 0); let manifest = quote_manifest(names); + let base_include = match kind { + ObjectKind::Archetype => "from .._baseclasses import Archetype", + ObjectKind::Component => "from .._baseclasses import Component", + ObjectKind::Datatype => "", + }; code.push_unindented_text( format!( " @@ -185,9 +194,11 @@ fn quote_objects( import numpy.typing as npt import pyarrow as pa - from dataclasses import dataclass + from dataclasses import dataclass, field from typing import Any, Dict, Iterable, Optional, Sequence, Set, Tuple, Union + {base_include} + __all__ = [{manifest}] ", @@ -220,14 +231,21 @@ fn quote_objects( let manifest = quote_manifest(mods.iter().flat_map(|(_, names)| names.iter())); + let (base_manifest, base_include) = match kind { + ObjectKind::Archetype => ("\"Archetype\", ", "from .._baseclasses import Archetype\n"), + ObjectKind::Component => ("\"Component\", ", "from .._baseclasses import Component\n"), + ObjectKind::Datatype => ("", ""), + }; + code.push_text(&format!("# {AUTOGEN_WARNING}"), 2, 0); code.push_unindented_text( format!( " from __future__ import annotations - __all__ = [{manifest}] + __all__ = [{base_manifest}{manifest}] + {base_include} ", ), 0, @@ -275,6 +293,10 @@ impl QuotedObject { let mut code = String::new(); + let superclass = match *kind { + ObjectKind::Archetype => "(Archetype)", + ObjectKind::Component | ObjectKind::Datatype => "", + }; code.push_unindented_text( format!( r#" @@ -282,7 +304,7 @@ impl QuotedObject { ## --- {name} --- ## @dataclass - class {name}: + class {name}{superclass}: "# ), 0, @@ -317,10 +339,18 @@ impl QuotedObject { } else { typ }; - let typ = if *is_nullable { - format!("{typ} | None = None") - } else { + let typ = if *kind == ObjectKind::Archetype { + if !*is_nullable { + format!("{typ} = field(metadata={{'component': 'primary'}})") + } else { + format!( + "{typ} | None = field(default=None, metadata={{'component': 'secondary'}})" + ) + } + } else if !*is_nullable { typ + } else { + format!("{typ} | None = None") }; code.push_text(format!("{name}: {typ}"), 1, 4); @@ -457,13 +487,13 @@ fn quote_str_repr_from_obj(obj: &Object) -> String { s = f"rr.{type(self).__name__}(\n" from dataclasses import fields - for field in fields(self): - data = getattr(self, field.name) - datatype = getattr(data, "type", None) - if datatype: - name = datatype.extension_name - typ = datatype.storage_type - s += f" {name}<{typ}>(\n {data.to_pylist()}\n )\n" + for fld in fields(self): + if "component" in fld.metadata: + comp: components.Component = getattr(self, fld.name) + if datatype := getattr(comp, "type"): + name = comp.extension_name + typ = datatype.storage_type + s += f" {name}<{typ}>(\n {comp.to_pylist()}\n )\n" s += ")" @@ -754,6 +784,12 @@ fn quote_arrow_support_from_obj(arrow_registry: &ArrowRegistry, obj: &Object) -> .try_get_attr::(ATTR_RERUN_LEGACY_FQNAME) .unwrap_or_else(|| fqname.clone()); + let superclass = if kind == &ObjectKind::Component { + "Component, " + } else { + "" + }; + unindent::unindent(&format!( r#" @@ -784,7 +820,9 @@ fn quote_arrow_support_from_obj(arrow_registry: &ArrowRegistry, obj: &Object) -> # TODO(cmc): bring back registration to pyarrow once legacy types are gone # pa.register_extension_type({arrow}()) - class {many}(pa.ExtensionArray, {many}Ext): # type: ignore[misc] + class {many}({superclass}{many}Ext): # type: ignore[misc] + _extension_name = "{legacy_fqname}" + @staticmethod def from_similar(data: {many_aliases} | None) -> pa.Array: if data is None: diff --git a/examples/python/api_demo/main.py b/examples/python/api_demo/main.py index 6fa22c021ef2..578df1e1cc8c 100755 --- a/examples/python/api_demo/main.py +++ b/examples/python/api_demo/main.py @@ -33,19 +33,36 @@ def run_segmentation() -> None: rr.log_segmentation_image("seg_demo/img", segmentation_img) # Log a bunch of classified 2D points - rr.log_point("seg_demo/single_point", np.array([64, 64]), class_id=13) - rr.log_point("seg_demo/single_point_labeled", np.array([90, 50]), class_id=13, label="labeled point") - rr.log_points("seg_demo/several_points0", np.array([[20, 50], [100, 70], [60, 30]]), class_ids=42) - rr.log_points( - "seg_demo/several_points1", - np.array([[40, 50], [120, 70], [80, 30]]), - class_ids=np.array([13, 42, 99], dtype=np.uint8), - ) - rr.log_points( - "seg_demo/many points", - np.array([[100 + (int(i / 5)) * 2, 100 + (i % 5) * 2] for i in range(25)]), - class_ids=np.array([42], dtype=np.uint8), - ) + if rr.ENABLE_NEXT_GEN_API: + # Note: this uses the new, WIP object-oriented API + rr.log_any("seg_demo/single_point", rr.Points2D([64, 64], class_ids=13)) + rr.log_any("seg_demo/single_point_labeled", rr.Points2D([90, 50], class_ids=13, labels="labeled point")) + rr.log_any("seg_demo/several_points0", rr.Points2D([[20, 50], [100, 70], [60, 30]], class_ids=42)) + rr.log_any( + "seg_demo/several_points1", + rr.Points2D([[40, 50], [120, 70], [80, 30]], class_ids=np.array([13, 42, 99], dtype=np.uint8)), + ) + rr.log_any( + "seg_demo/many points", + rr.Points2D( + [[100 + (int(i / 5)) * 2, 100 + (i % 5) * 2] for i in range(25)], + class_ids=np.array([42], dtype=np.uint8), + ), + ) + else: + rr.log_point("seg_demo/single_point", np.array([64, 64]), class_id=13) + rr.log_point("seg_demo/single_point_labeled", np.array([90, 50]), class_id=13, label="labeled point") + rr.log_points("seg_demo/several_points0", np.array([[20, 50], [100, 70], [60, 30]]), class_ids=42) + rr.log_points( + "seg_demo/several_points1", + np.array([[40, 50], [120, 70], [80, 30]]), + class_ids=np.array([13, 42, 99], dtype=np.uint8), + ) + rr.log_points( + "seg_demo/many points", + np.array([[100 + (int(i / 5)) * 2, 100 + (i % 5) * 2] for i in range(25)]), + class_ids=np.array([42], dtype=np.uint8), + ) rr.log_text_entry("logs/seg_demo_log", "default colored rects, default colored points, a single point has a label") diff --git a/rerun_py/rerun_sdk/rerun/__init__.py b/rerun_py/rerun_sdk/rerun/__init__.py index b318cfea0c7b..64184c10318c 100644 --- a/rerun_py/rerun_sdk/rerun/__init__.py +++ b/rerun_py/rerun_sdk/rerun/__init__.py @@ -56,10 +56,10 @@ from .time import reset_time, set_time_nanos, set_time_seconds, set_time_sequence # Next-gen API imports -# TODO(ab): remove this guard, here to make it easy to "hide" the next gen API if needed in the short term. -_ENABLE_NEXT_GEN_API = True -if _ENABLE_NEXT_GEN_API: +ENABLE_NEXT_GEN_API = True +if ENABLE_NEXT_GEN_API: from ._rerun2.archetypes import * + from ._rerun2.log_any import log_any def _init_recording_stream() -> None: diff --git a/rerun_py/rerun_sdk/rerun/_rerun2/_baseclasses.py b/rerun_py/rerun_sdk/rerun/_rerun2/_baseclasses.py new file mode 100644 index 000000000000..a42f49a61ad3 --- /dev/null +++ b/rerun_py/rerun_sdk/rerun/_rerun2/_baseclasses.py @@ -0,0 +1,16 @@ +from __future__ import annotations + +from dataclasses import dataclass + +import pyarrow as pa + + +@dataclass +class Archetype: + pass + + +class Component(pa.ExtensionArray): # type: ignore[misc] + @property + def extension_name(self) -> str: + return getattr(self, "_extension_name", "") diff --git a/rerun_py/rerun_sdk/rerun/_rerun2/archetypes/__init__.py b/rerun_py/rerun_sdk/rerun/_rerun2/archetypes/__init__.py index b63e459e6d55..b718bbc1bf28 100644 --- a/rerun_py/rerun_sdk/rerun/_rerun2/archetypes/__init__.py +++ b/rerun_py/rerun_sdk/rerun/_rerun2/archetypes/__init__.py @@ -2,7 +2,8 @@ from __future__ import annotations -__all__ = ["AffixFuzzer1", "Points2D"] +__all__ = ["Archetype", "AffixFuzzer1", "Points2D"] +from .._baseclasses import Archetype from .fuzzy import AffixFuzzer1 from .points2d import Points2D diff --git a/rerun_py/rerun_sdk/rerun/_rerun2/archetypes/fuzzy.py b/rerun_py/rerun_sdk/rerun/_rerun2/archetypes/fuzzy.py index fcff0b4adb97..131509f65096 100644 --- a/rerun_py/rerun_sdk/rerun/_rerun2/archetypes/fuzzy.py +++ b/rerun_py/rerun_sdk/rerun/_rerun2/archetypes/fuzzy.py @@ -2,7 +2,9 @@ from __future__ import annotations -from dataclasses import dataclass +from dataclasses import dataclass, field + +from .._baseclasses import Archetype __all__ = ["AffixFuzzer1"] @@ -12,48 +14,48 @@ @dataclass -class AffixFuzzer1: - fuzz1001: components.AffixFuzzer1Array - fuzz1002: components.AffixFuzzer2Array - fuzz1003: components.AffixFuzzer3Array - fuzz1004: components.AffixFuzzer4Array - fuzz1005: components.AffixFuzzer5Array - fuzz1006: components.AffixFuzzer6Array - fuzz1007: components.AffixFuzzer7Array - fuzz1101: components.AffixFuzzer1Array - fuzz1102: components.AffixFuzzer2Array - fuzz1103: components.AffixFuzzer3Array - fuzz1104: components.AffixFuzzer4Array - fuzz1105: components.AffixFuzzer5Array - fuzz1106: components.AffixFuzzer6Array - fuzz1107: components.AffixFuzzer7Array - fuzz2001: components.AffixFuzzer1Array | None = None - fuzz2002: components.AffixFuzzer2Array | None = None - fuzz2003: components.AffixFuzzer3Array | None = None - fuzz2004: components.AffixFuzzer4Array | None = None - fuzz2005: components.AffixFuzzer5Array | None = None - fuzz2006: components.AffixFuzzer6Array | None = None - fuzz2007: components.AffixFuzzer7Array | None = None - fuzz2101: components.AffixFuzzer1Array | None = None - fuzz2102: components.AffixFuzzer2Array | None = None - fuzz2103: components.AffixFuzzer3Array | None = None - fuzz2104: components.AffixFuzzer4Array | None = None - fuzz2105: components.AffixFuzzer5Array | None = None - fuzz2106: components.AffixFuzzer6Array | None = None - fuzz2107: components.AffixFuzzer7Array | None = None +class AffixFuzzer1(Archetype): + fuzz1001: components.AffixFuzzer1Array = field(metadata={"component": "primary"}) + fuzz1002: components.AffixFuzzer2Array = field(metadata={"component": "primary"}) + fuzz1003: components.AffixFuzzer3Array = field(metadata={"component": "primary"}) + fuzz1004: components.AffixFuzzer4Array = field(metadata={"component": "primary"}) + fuzz1005: components.AffixFuzzer5Array = field(metadata={"component": "primary"}) + fuzz1006: components.AffixFuzzer6Array = field(metadata={"component": "primary"}) + fuzz1007: components.AffixFuzzer7Array = field(metadata={"component": "primary"}) + fuzz1101: components.AffixFuzzer1Array = field(metadata={"component": "primary"}) + fuzz1102: components.AffixFuzzer2Array = field(metadata={"component": "primary"}) + fuzz1103: components.AffixFuzzer3Array = field(metadata={"component": "primary"}) + fuzz1104: components.AffixFuzzer4Array = field(metadata={"component": "primary"}) + fuzz1105: components.AffixFuzzer5Array = field(metadata={"component": "primary"}) + fuzz1106: components.AffixFuzzer6Array = field(metadata={"component": "primary"}) + fuzz1107: components.AffixFuzzer7Array = field(metadata={"component": "primary"}) + fuzz2001: components.AffixFuzzer1Array | None = field(default=None, metadata={"component": "secondary"}) + fuzz2002: components.AffixFuzzer2Array | None = field(default=None, metadata={"component": "secondary"}) + fuzz2003: components.AffixFuzzer3Array | None = field(default=None, metadata={"component": "secondary"}) + fuzz2004: components.AffixFuzzer4Array | None = field(default=None, metadata={"component": "secondary"}) + fuzz2005: components.AffixFuzzer5Array | None = field(default=None, metadata={"component": "secondary"}) + fuzz2006: components.AffixFuzzer6Array | None = field(default=None, metadata={"component": "secondary"}) + fuzz2007: components.AffixFuzzer7Array | None = field(default=None, metadata={"component": "secondary"}) + fuzz2101: components.AffixFuzzer1Array | None = field(default=None, metadata={"component": "secondary"}) + fuzz2102: components.AffixFuzzer2Array | None = field(default=None, metadata={"component": "secondary"}) + fuzz2103: components.AffixFuzzer3Array | None = field(default=None, metadata={"component": "secondary"}) + fuzz2104: components.AffixFuzzer4Array | None = field(default=None, metadata={"component": "secondary"}) + fuzz2105: components.AffixFuzzer5Array | None = field(default=None, metadata={"component": "secondary"}) + fuzz2106: components.AffixFuzzer6Array | None = field(default=None, metadata={"component": "secondary"}) + fuzz2107: components.AffixFuzzer7Array | None = field(default=None, metadata={"component": "secondary"}) def __str__(self) -> str: s = f"rr.{type(self).__name__}(\n" from dataclasses import fields - for field in fields(self): - data = getattr(self, field.name) - datatype = getattr(data, "type", None) - if datatype: - name = datatype.extension_name - typ = datatype.storage_type - s += f" {name}<{typ}>(\n {data.to_pylist()}\n )\n" + for fld in fields(self): + if "component" in fld.metadata: + comp: components.Component = getattr(self, fld.name) + if datatype := getattr(comp, "type"): + name = comp.extension_name + typ = datatype.storage_type + s += f" {name}<{typ}>(\n {comp.to_pylist()}\n )\n" s += ")" diff --git a/rerun_py/rerun_sdk/rerun/_rerun2/archetypes/points2d.py b/rerun_py/rerun_sdk/rerun/_rerun2/archetypes/points2d.py index 6fd3d30c8795..137fefc78fd3 100644 --- a/rerun_py/rerun_sdk/rerun/_rerun2/archetypes/points2d.py +++ b/rerun_py/rerun_sdk/rerun/_rerun2/archetypes/points2d.py @@ -2,7 +2,9 @@ from __future__ import annotations -from dataclasses import dataclass +from dataclasses import dataclass, field + +from .._baseclasses import Archetype __all__ = ["Points2D"] @@ -12,20 +14,20 @@ @dataclass -class Points2D: +class Points2D(Archetype): """A 2D point cloud with positions and optional colors, radii, labels, etc.""" - points: components.Point2DArray + points: components.Point2DArray = field(metadata={"component": "primary"}) """ All the actual 2D points that make up the point cloud. """ - radii: components.RadiusArray | None = None + radii: components.RadiusArray | None = field(default=None, metadata={"component": "secondary"}) """ Optional radii for the points, effectively turning them into circles. """ - colors: components.ColorArray | None = None + colors: components.ColorArray | None = field(default=None, metadata={"component": "secondary"}) """ Optional colors for the points. @@ -33,12 +35,12 @@ class Points2D: As either 0-1 floats or 0-255 integers, with separate alpha. """ - labels: components.LabelArray | None = None + labels: components.LabelArray | None = field(default=None, metadata={"component": "secondary"}) """ Optional text labels for the points. """ - draw_order: components.DrawOrderArray | None = None + draw_order: components.DrawOrderArray | None = field(default=None, metadata={"component": "secondary"}) """ An optional floating point value that specifies the 2D drawing order. Objects with higher values are drawn on top of those with lower values. @@ -46,14 +48,14 @@ class Points2D: The default for 2D points is 30.0. """ - class_ids: components.ClassIdArray | None = None + class_ids: components.ClassIdArray | None = field(default=None, metadata={"component": "secondary"}) """ Optional class Ids for the points. The class ID provides colors and labels if not specified explicitly. """ - keypoint_ids: components.KeypointIdArray | None = None + keypoint_ids: components.KeypointIdArray | None = field(default=None, metadata={"component": "secondary"}) """ Optional keypoint IDs for the points, identifying them within a class. @@ -65,7 +67,7 @@ class Points2D: detected skeleton. """ - instance_keys: components.InstanceKeyArray | None = None + instance_keys: components.InstanceKeyArray | None = field(default=None, metadata={"component": "secondary"}) """ Unique identifiers for each individual point in the batch. """ @@ -75,13 +77,13 @@ def __str__(self) -> str: from dataclasses import fields - for field in fields(self): - data = getattr(self, field.name) - datatype = getattr(data, "type", None) - if datatype: - name = datatype.extension_name - typ = datatype.storage_type - s += f" {name}<{typ}>(\n {data.to_pylist()}\n )\n" + for fld in fields(self): + if "component" in fld.metadata: + comp: components.Component = getattr(self, fld.name) + if datatype := getattr(comp, "type"): + name = comp.extension_name + typ = datatype.storage_type + s += f" {name}<{typ}>(\n {comp.to_pylist()}\n )\n" s += ")" diff --git a/rerun_py/rerun_sdk/rerun/_rerun2/components/__init__.py b/rerun_py/rerun_sdk/rerun/_rerun2/components/__init__.py index 39cab6189458..544020ec3748 100644 --- a/rerun_py/rerun_sdk/rerun/_rerun2/components/__init__.py +++ b/rerun_py/rerun_sdk/rerun/_rerun2/components/__init__.py @@ -3,6 +3,7 @@ from __future__ import annotations __all__ = [ + "Component", "AffixFuzzer1", "AffixFuzzer1Array", "AffixFuzzer1ArrayLike", @@ -80,6 +81,7 @@ "RadiusType", ] +from .._baseclasses import Component from .class_id import ClassId, ClassIdArray, ClassIdArrayLike, ClassIdLike, ClassIdType from .color import Color, ColorArray, ColorArrayLike, ColorLike, ColorType from .draw_order import DrawOrder, DrawOrderArray, DrawOrderArrayLike, DrawOrderLike, DrawOrderType diff --git a/rerun_py/rerun_sdk/rerun/_rerun2/components/class_id.py b/rerun_py/rerun_sdk/rerun/_rerun2/components/class_id.py index 0cd3d86043c4..9b62102c302f 100644 --- a/rerun_py/rerun_sdk/rerun/_rerun2/components/class_id.py +++ b/rerun_py/rerun_sdk/rerun/_rerun2/components/class_id.py @@ -9,6 +9,8 @@ import numpy.typing as npt import pyarrow as pa +from .._baseclasses import Component + __all__ = ["ClassId", "ClassIdArray", "ClassIdArrayLike", "ClassIdLike", "ClassIdType"] @@ -65,7 +67,9 @@ def __arrow_ext_class__(self: type[pa.ExtensionType]) -> type[pa.ExtensionArray] # pa.register_extension_type(ClassIdType()) -class ClassIdArray(pa.ExtensionArray, ClassIdArrayExt): # type: ignore[misc] +class ClassIdArray(Component, ClassIdArrayExt): # type: ignore[misc] + _extension_name = "rerun.class_id" + @staticmethod def from_similar(data: ClassIdArrayLike | None) -> pa.Array: if data is None: diff --git a/rerun_py/rerun_sdk/rerun/_rerun2/components/color.py b/rerun_py/rerun_sdk/rerun/_rerun2/components/color.py index 327bf8bee058..76fc3e37c376 100644 --- a/rerun_py/rerun_sdk/rerun/_rerun2/components/color.py +++ b/rerun_py/rerun_sdk/rerun/_rerun2/components/color.py @@ -9,6 +9,8 @@ import numpy.typing as npt import pyarrow as pa +from .._baseclasses import Component + __all__ = ["Color", "ColorArray", "ColorArrayLike", "ColorLike", "ColorType"] @@ -74,7 +76,9 @@ def __arrow_ext_class__(self: type[pa.ExtensionType]) -> type[pa.ExtensionArray] # pa.register_extension_type(ColorType()) -class ColorArray(pa.ExtensionArray, ColorArrayExt): # type: ignore[misc] +class ColorArray(Component, ColorArrayExt): # type: ignore[misc] + _extension_name = "rerun.colorrgba" + @staticmethod def from_similar(data: ColorArrayLike | None) -> pa.Array: if data is None: diff --git a/rerun_py/rerun_sdk/rerun/_rerun2/components/draw_order.py b/rerun_py/rerun_sdk/rerun/_rerun2/components/draw_order.py index b015c66feb4f..59335a0c4143 100644 --- a/rerun_py/rerun_sdk/rerun/_rerun2/components/draw_order.py +++ b/rerun_py/rerun_sdk/rerun/_rerun2/components/draw_order.py @@ -9,6 +9,8 @@ import numpy.typing as npt import pyarrow as pa +from .._baseclasses import Component + __all__ = ["DrawOrder", "DrawOrderArray", "DrawOrderArrayLike", "DrawOrderLike", "DrawOrderType"] @@ -66,7 +68,9 @@ def __arrow_ext_class__(self: type[pa.ExtensionType]) -> type[pa.ExtensionArray] # pa.register_extension_type(DrawOrderType()) -class DrawOrderArray(pa.ExtensionArray, DrawOrderArrayExt): # type: ignore[misc] +class DrawOrderArray(Component, DrawOrderArrayExt): # type: ignore[misc] + _extension_name = "rerun.draw_order" + @staticmethod def from_similar(data: DrawOrderArrayLike | None) -> pa.Array: if data is None: diff --git a/rerun_py/rerun_sdk/rerun/_rerun2/components/fuzzy.py b/rerun_py/rerun_sdk/rerun/_rerun2/components/fuzzy.py index 905feacc6a6e..24e4de7d4737 100644 --- a/rerun_py/rerun_sdk/rerun/_rerun2/components/fuzzy.py +++ b/rerun_py/rerun_sdk/rerun/_rerun2/components/fuzzy.py @@ -8,6 +8,8 @@ import numpy.typing as npt import pyarrow as pa +from .._baseclasses import Component + __all__ = [ "AffixFuzzer1", "AffixFuzzer1Array", @@ -104,7 +106,9 @@ def __arrow_ext_class__(self: type[pa.ExtensionType]) -> type[pa.ExtensionArray] # pa.register_extension_type(AffixFuzzer1Type()) -class AffixFuzzer1Array(pa.ExtensionArray, AffixFuzzer1ArrayExt): # type: ignore[misc] +class AffixFuzzer1Array(Component, AffixFuzzer1ArrayExt): # type: ignore[misc] + _extension_name = "rerun.testing.components.AffixFuzzer1" + @staticmethod def from_similar(data: AffixFuzzer1ArrayLike | None) -> pa.Array: if data is None: @@ -176,7 +180,9 @@ def __arrow_ext_class__(self: type[pa.ExtensionType]) -> type[pa.ExtensionArray] # pa.register_extension_type(AffixFuzzer2Type()) -class AffixFuzzer2Array(pa.ExtensionArray, AffixFuzzer2ArrayExt): # type: ignore[misc] +class AffixFuzzer2Array(Component, AffixFuzzer2ArrayExt): # type: ignore[misc] + _extension_name = "rerun.testing.components.AffixFuzzer2" + @staticmethod def from_similar(data: AffixFuzzer2ArrayLike | None) -> pa.Array: if data is None: @@ -263,7 +269,9 @@ def __arrow_ext_class__(self: type[pa.ExtensionType]) -> type[pa.ExtensionArray] # pa.register_extension_type(AffixFuzzer3Type()) -class AffixFuzzer3Array(pa.ExtensionArray, AffixFuzzer3ArrayExt): # type: ignore[misc] +class AffixFuzzer3Array(Component, AffixFuzzer3ArrayExt): # type: ignore[misc] + _extension_name = "rerun.testing.components.AffixFuzzer3" + @staticmethod def from_similar(data: AffixFuzzer3ArrayLike | None) -> pa.Array: if data is None: @@ -335,7 +343,9 @@ def __arrow_ext_class__(self: type[pa.ExtensionType]) -> type[pa.ExtensionArray] # pa.register_extension_type(AffixFuzzer4Type()) -class AffixFuzzer4Array(pa.ExtensionArray, AffixFuzzer4ArrayExt): # type: ignore[misc] +class AffixFuzzer4Array(Component, AffixFuzzer4ArrayExt): # type: ignore[misc] + _extension_name = "rerun.testing.components.AffixFuzzer4" + @staticmethod def from_similar(data: AffixFuzzer4ArrayLike | None) -> pa.Array: if data is None: @@ -407,7 +417,9 @@ def __arrow_ext_class__(self: type[pa.ExtensionType]) -> type[pa.ExtensionArray] # pa.register_extension_type(AffixFuzzer5Type()) -class AffixFuzzer5Array(pa.ExtensionArray, AffixFuzzer5ArrayExt): # type: ignore[misc] +class AffixFuzzer5Array(Component, AffixFuzzer5ArrayExt): # type: ignore[misc] + _extension_name = "rerun.testing.components.AffixFuzzer5" + @staticmethod def from_similar(data: AffixFuzzer5ArrayLike | None) -> pa.Array: if data is None: @@ -494,7 +506,9 @@ def __arrow_ext_class__(self: type[pa.ExtensionType]) -> type[pa.ExtensionArray] # pa.register_extension_type(AffixFuzzer6Type()) -class AffixFuzzer6Array(pa.ExtensionArray, AffixFuzzer6ArrayExt): # type: ignore[misc] +class AffixFuzzer6Array(Component, AffixFuzzer6ArrayExt): # type: ignore[misc] + _extension_name = "rerun.testing.components.AffixFuzzer6" + @staticmethod def from_similar(data: AffixFuzzer6ArrayLike | None) -> pa.Array: if data is None: @@ -609,7 +623,9 @@ def __arrow_ext_class__(self: type[pa.ExtensionType]) -> type[pa.ExtensionArray] # pa.register_extension_type(AffixFuzzer7Type()) -class AffixFuzzer7Array(pa.ExtensionArray, AffixFuzzer7ArrayExt): # type: ignore[misc] +class AffixFuzzer7Array(Component, AffixFuzzer7ArrayExt): # type: ignore[misc] + _extension_name = "rerun.testing.components.AffixFuzzer7" + @staticmethod def from_similar(data: AffixFuzzer7ArrayLike | None) -> pa.Array: if data is None: diff --git a/rerun_py/rerun_sdk/rerun/_rerun2/components/instance_key.py b/rerun_py/rerun_sdk/rerun/_rerun2/components/instance_key.py index 01b825e7beb5..10a85e14daa1 100644 --- a/rerun_py/rerun_sdk/rerun/_rerun2/components/instance_key.py +++ b/rerun_py/rerun_sdk/rerun/_rerun2/components/instance_key.py @@ -9,6 +9,8 @@ import numpy.typing as npt import pyarrow as pa +from .._baseclasses import Component + __all__ = ["InstanceKey", "InstanceKeyArray", "InstanceKeyArrayLike", "InstanceKeyLike", "InstanceKeyType"] @@ -58,7 +60,9 @@ def __arrow_ext_class__(self: type[pa.ExtensionType]) -> type[pa.ExtensionArray] # pa.register_extension_type(InstanceKeyType()) -class InstanceKeyArray(pa.ExtensionArray, InstanceKeyArrayExt): # type: ignore[misc] +class InstanceKeyArray(Component, InstanceKeyArrayExt): # type: ignore[misc] + _extension_name = "rerun.instance_key" + @staticmethod def from_similar(data: InstanceKeyArrayLike | None) -> pa.Array: if data is None: diff --git a/rerun_py/rerun_sdk/rerun/_rerun2/components/instance_key_ext.py b/rerun_py/rerun_sdk/rerun/_rerun2/components/instance_key_ext.py index 79a98ffedef3..55fdf14f9981 100644 --- a/rerun_py/rerun_sdk/rerun/_rerun2/components/instance_key_ext.py +++ b/rerun_py/rerun_sdk/rerun/_rerun2/components/instance_key_ext.py @@ -2,11 +2,16 @@ __all__ = ["InstanceKeyArrayExt"] -from typing import Any, Sequence +from typing import TYPE_CHECKING, Any, Sequence import numpy as np import pyarrow as pa +if TYPE_CHECKING: + from .instance_key import InstanceKeyArray + +_MAX_U64 = 2**64 - 1 + class InstanceKeyArrayExt: @staticmethod @@ -19,3 +24,10 @@ def _from_similar( array = np.asarray(data, dtype=np.uint64).flatten() return arrow().wrap_array(pa.array(array, type=arrow().storage_type)) + + @staticmethod + def splat() -> InstanceKeyArray: + from .instance_key import InstanceKeyType + + storage = pa.array([_MAX_U64], type=InstanceKeyType().storage_type) + return storage # type: ignore[no-any-return] diff --git a/rerun_py/rerun_sdk/rerun/_rerun2/components/keypoint_id.py b/rerun_py/rerun_sdk/rerun/_rerun2/components/keypoint_id.py index 38f9b3ddbac1..e53b0bcb6764 100644 --- a/rerun_py/rerun_sdk/rerun/_rerun2/components/keypoint_id.py +++ b/rerun_py/rerun_sdk/rerun/_rerun2/components/keypoint_id.py @@ -9,6 +9,8 @@ import numpy.typing as npt import pyarrow as pa +from .._baseclasses import Component + __all__ = ["KeypointId", "KeypointIdArray", "KeypointIdArrayLike", "KeypointIdLike", "KeypointIdType"] @@ -67,7 +69,9 @@ def __arrow_ext_class__(self: type[pa.ExtensionType]) -> type[pa.ExtensionArray] # pa.register_extension_type(KeypointIdType()) -class KeypointIdArray(pa.ExtensionArray, KeypointIdArrayExt): # type: ignore[misc] +class KeypointIdArray(Component, KeypointIdArrayExt): # type: ignore[misc] + _extension_name = "rerun.keypoint_id" + @staticmethod def from_similar(data: KeypointIdArrayLike | None) -> pa.Array: if data is None: diff --git a/rerun_py/rerun_sdk/rerun/_rerun2/components/label.py b/rerun_py/rerun_sdk/rerun/_rerun2/components/label.py index 5af6ea0c614b..b88d61613761 100644 --- a/rerun_py/rerun_sdk/rerun/_rerun2/components/label.py +++ b/rerun_py/rerun_sdk/rerun/_rerun2/components/label.py @@ -7,6 +7,8 @@ import pyarrow as pa +from .._baseclasses import Component + __all__ = ["Label", "LabelArray", "LabelArrayLike", "LabelLike", "LabelType"] @@ -59,7 +61,9 @@ def __arrow_ext_class__(self: type[pa.ExtensionType]) -> type[pa.ExtensionArray] # pa.register_extension_type(LabelType()) -class LabelArray(pa.ExtensionArray, LabelArrayExt): # type: ignore[misc] +class LabelArray(Component, LabelArrayExt): # type: ignore[misc] + _extension_name = "rerun.label" + @staticmethod def from_similar(data: LabelArrayLike | None) -> pa.Array: if data is None: diff --git a/rerun_py/rerun_sdk/rerun/_rerun2/components/label_ext.py b/rerun_py/rerun_sdk/rerun/_rerun2/components/label_ext.py index 5baecff2fed7..1a74aef9bb08 100644 --- a/rerun_py/rerun_sdk/rerun/_rerun2/components/label_ext.py +++ b/rerun_py/rerun_sdk/rerun/_rerun2/components/label_ext.py @@ -12,7 +12,9 @@ class LabelArrayExt: def _from_similar( data: Any | None, *, mono: type, mono_aliases: Any, many: type, many_aliases: Any, arrow: type ) -> pa.Array: - if isinstance(data, Sequence): + if isinstance(data, str): + array = [data] + elif isinstance(data, Sequence): array = [str(datum) for datum in data] else: array = [str(data)] diff --git a/rerun_py/rerun_sdk/rerun/_rerun2/components/point2d.py b/rerun_py/rerun_sdk/rerun/_rerun2/components/point2d.py index a7d5aaac503e..e456b9791440 100644 --- a/rerun_py/rerun_sdk/rerun/_rerun2/components/point2d.py +++ b/rerun_py/rerun_sdk/rerun/_rerun2/components/point2d.py @@ -9,6 +9,8 @@ import numpy.typing as npt import pyarrow as pa +from .._baseclasses import Component + __all__ = ["Point2D", "Point2DArray", "Point2DArrayLike", "Point2DLike", "Point2DType"] @@ -60,7 +62,9 @@ def __arrow_ext_class__(self: type[pa.ExtensionType]) -> type[pa.ExtensionArray] # pa.register_extension_type(Point2DType()) -class Point2DArray(pa.ExtensionArray, Point2DArrayExt): # type: ignore[misc] +class Point2DArray(Component, Point2DArrayExt): # type: ignore[misc] + _extension_name = "rerun.point2d" + @staticmethod def from_similar(data: Point2DArrayLike | None) -> pa.Array: if data is None: diff --git a/rerun_py/rerun_sdk/rerun/_rerun2/components/radius.py b/rerun_py/rerun_sdk/rerun/_rerun2/components/radius.py index ec81b9eeffe9..407755538e74 100644 --- a/rerun_py/rerun_sdk/rerun/_rerun2/components/radius.py +++ b/rerun_py/rerun_sdk/rerun/_rerun2/components/radius.py @@ -9,6 +9,8 @@ import numpy.typing as npt import pyarrow as pa +from .._baseclasses import Component + __all__ = ["Radius", "RadiusArray", "RadiusArrayLike", "RadiusLike", "RadiusType"] @@ -58,7 +60,9 @@ def __arrow_ext_class__(self: type[pa.ExtensionType]) -> type[pa.ExtensionArray] # pa.register_extension_type(RadiusType()) -class RadiusArray(pa.ExtensionArray, RadiusArrayExt): # type: ignore[misc] +class RadiusArray(Component, RadiusArrayExt): # type: ignore[misc] + _extension_name = "rerun.radius" + @staticmethod def from_similar(data: RadiusArrayLike | None) -> pa.Array: if data is None: diff --git a/rerun_py/rerun_sdk/rerun/_rerun2/datatypes/__init__.py b/rerun_py/rerun_sdk/rerun/_rerun2/datatypes/__init__.py index 9519be379256..cec81dfa19cf 100644 --- a/rerun_py/rerun_sdk/rerun/_rerun2/datatypes/__init__.py +++ b/rerun_py/rerun_sdk/rerun/_rerun2/datatypes/__init__.py @@ -20,6 +20,7 @@ "Vec2DType", ] + from .fuzzy import ( AffixFuzzer1, AffixFuzzer1Array, diff --git a/rerun_py/rerun_sdk/rerun/_rerun2/datatypes/fuzzy.py b/rerun_py/rerun_sdk/rerun/_rerun2/datatypes/fuzzy.py index 76d2ee4da208..f095027ea27e 100644 --- a/rerun_py/rerun_sdk/rerun/_rerun2/datatypes/fuzzy.py +++ b/rerun_py/rerun_sdk/rerun/_rerun2/datatypes/fuzzy.py @@ -84,7 +84,9 @@ def __arrow_ext_class__(self: type[pa.ExtensionType]) -> type[pa.ExtensionArray] # pa.register_extension_type(AffixFuzzer1Type()) -class AffixFuzzer1Array(pa.ExtensionArray, AffixFuzzer1ArrayExt): # type: ignore[misc] +class AffixFuzzer1Array(AffixFuzzer1ArrayExt): # type: ignore[misc] + _extension_name = "rerun.testing.datatypes.AffixFuzzer1" + @staticmethod def from_similar(data: AffixFuzzer1ArrayLike | None) -> pa.Array: if data is None: @@ -146,7 +148,9 @@ def __arrow_ext_class__(self: type[pa.ExtensionType]) -> type[pa.ExtensionArray] # pa.register_extension_type(AffixFuzzer2Type()) -class AffixFuzzer2Array(pa.ExtensionArray, AffixFuzzer2ArrayExt): # type: ignore[misc] +class AffixFuzzer2Array(AffixFuzzer2ArrayExt): # type: ignore[misc] + _extension_name = "rerun.testing.datatypes.AffixFuzzer2" + @staticmethod def from_similar(data: AffixFuzzer2ArrayLike | None) -> pa.Array: if data is None: diff --git a/rerun_py/rerun_sdk/rerun/_rerun2/datatypes/vec2d.py b/rerun_py/rerun_sdk/rerun/_rerun2/datatypes/vec2d.py index 2efcf56fbcf2..14aa1d0a2a1c 100644 --- a/rerun_py/rerun_sdk/rerun/_rerun2/datatypes/vec2d.py +++ b/rerun_py/rerun_sdk/rerun/_rerun2/datatypes/vec2d.py @@ -60,7 +60,9 @@ def __arrow_ext_class__(self: type[pa.ExtensionType]) -> type[pa.ExtensionArray] # pa.register_extension_type(Vec2DType()) -class Vec2DArray(pa.ExtensionArray, Vec2DArrayExt): # type: ignore[misc] +class Vec2DArray(Vec2DArrayExt): # type: ignore[misc] + _extension_name = "rerun.datatypes.Vec2D" + @staticmethod def from_similar(data: Vec2DArrayLike | None) -> pa.Array: if data is None: diff --git a/rerun_py/rerun_sdk/rerun/_rerun2/log_any.py b/rerun_py/rerun_sdk/rerun/_rerun2/log_any.py new file mode 100644 index 000000000000..6d7c4bbd7bb0 --- /dev/null +++ b/rerun_py/rerun_sdk/rerun/_rerun2/log_any.py @@ -0,0 +1,130 @@ +from __future__ import annotations + +from dataclasses import fields +from typing import Any, Iterable + +import numpy as np +import numpy.typing as npt +import pyarrow as pa + +from .. import RecordingStream, bindings +from ..log import error_utils +from .archetypes import Archetype +from .components import Component, InstanceKeyArray + +__all__ = ["log_any"] + + +EXT_PREFIX = "ext." + +ext_component_types: dict[str, Any] = {} + + +# adapted from rerun.log._add_extension_components +def _add_extension_components( + instanced: dict[str, Component], + splats: dict[str, Component], + ext: dict[str, Any], + identifiers: npt.NDArray[np.uint64] | None, +) -> None: + global ext_component_types + + for name, value in ext.items(): + # Don't log empty components + if value is None: + continue + + # Add the ext prefix, unless it's already there + if not name.startswith(EXT_PREFIX): + name = EXT_PREFIX + name + + np_type, pa_type = ext_component_types.get(name, (None, None)) + + try: + if np_type is not None: + np_value = np.atleast_1d(np.array(value, copy=False, dtype=np_type)) + pa_value = pa.array(np_value, type=pa_type) + else: + np_value = np.atleast_1d(np.array(value, copy=False)) + pa_value = pa.array(np_value) + ext_component_types[name] = (np_value.dtype, pa_value.type) + except Exception as ex: + error_utils._send_warning( + "Error converting extension data to arrow for component {}. Dropping.\n{}: {}".format( + name, type(ex).__name__, ex + ), + 1, + ) + continue + + is_splat = (len(np_value) == 1) and (len(identifiers or []) != 1) + + if is_splat: + splats[name] = pa_value # noqa + else: + instanced[name] = pa_value # noqa + + +def _extract_components(entity: Archetype) -> Iterable[tuple[Component, bool]]: + """Extract the components from an entity, yielding (component, is_primary) tuples.""" + for fld in fields(entity): + if "component" in fld.metadata: + yield getattr(entity, fld.name), fld.metadata["component"] == "primary" + + +def log_any( + entity_path: str, + entity: Archetype, + ext: dict[str, Any] | None = None, + timeless: bool = False, + recording: RecordingStream | None = None, +) -> None: + """ + Log an entity. + + Parameters + ---------- + entity_path: + Path to the entity in the space hierarchy. + entity: Archetype + The archetype object representing the entity. + ext: + Optional dictionary of extension components. See [rerun.log_extension_components][] + timeless: + If true, the entity will be timeless (default: False). + recording: + Specifies the [`rerun.RecordingStream`][] to use. + If left unspecified, defaults to the current active data recording, if there is one. + See also: [`rerun.init`][], [`rerun.set_global_data_recording`][]. + + """ + + from .. import strict_mode + + if strict_mode(): + if not isinstance(entity, Archetype): + raise TypeError(f"Expected Archetype, got {type(entity)}") + + instanced: dict[str, Component] = {} + splats: dict[str, Component] = {} + + # find canonical length of this entity by based on the longest length of any primary component + archetype_length = max(len(comp) for comp, primary in _extract_components(entity) if primary) + + for comp, primary in _extract_components(entity): + if primary: + instanced[comp.extension_name] = comp.storage + elif len(comp) == 1 and archetype_length > 1: + splats[comp.extension_name] = comp.storage + elif len(comp) > 1: + instanced[comp.extension_name] = comp.storage + + if ext: + _add_extension_components(instanced, splats, ext, None) + + if splats: + splats["rerun.instance_key"] = InstanceKeyArray.splat() + bindings.log_arrow_msg(entity_path, components=splats, timeless=timeless, recording=recording) + + # Always the primary component last so range-based queries will include the other data. See(#1215) + bindings.log_arrow_msg(entity_path, components=instanced, timeless=timeless, recording=recording) diff --git a/rerun_py/tests/unit/test_log.py b/rerun_py/tests/unit/test_log.py new file mode 100644 index 000000000000..b41b833df6ef --- /dev/null +++ b/rerun_py/tests/unit/test_log.py @@ -0,0 +1,10 @@ +from __future__ import annotations + +import rerun as rr + + +def test_log_point2d_basic() -> None: + """Basic test: logging a point shouldn't raise an exception...""" + points = rr.Points2D([(0, 0), (2, 2), (2, 2.5), (2.5, 2), (3, 4)], radii=0.5) + rr.init("test_log") + rr.log_any("points", points)