Python backport: add log_any()

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.
-128cd542f3f9e6cf3b2a44aeb022cd5f3ad819b00ce5371eeb311300f3e9a7f1
+c0d06250415f4b910072a218e82b969fb33ebca44ef59d544373439d6bb135f1

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<Path>,
     arrow_registry: &ArrowRegistry,
     all_objects: &Objects,
+    kind: ObjectKind,
     objs: &[&Object],
 ) -> (Vec<PathBuf>, Vec<String>) {
     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,14 +293,18 @@ impl QuotedObject {
 
         let mut code = String::new();
 
+        let superclass = match *kind {
+            ObjectKind::Archetype => "(Archetype)",
+            ObjectKind::Component | ObjectKind::Datatype => "",
+        };
         code.push_unindented_text(
             format!(
                 r#"
 
                 ## --- {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::<String>(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:

examples/python/api_demo/main.py

@@ -33,19 +33,38 @@ 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(np.array([64, 64]), class_ids=13))
+        rr.log_any(
+            "seg_demo/single_point_labeled", rr.Points2D(np.array([90, 50]), class_ids=13, labels="labeled point")
+        )
+        rr.log_any("seg_demo/several_points0", rr.Points2D(np.array([[20, 50], [100, 70], [60, 30]]), class_ids=42))
+        rr.log_any(
+            "seg_demo/several_points1",
+            rr.Points2D(np.array([[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(
+                np.array([[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")
 

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:

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

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

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