Codegen/IDL 7: handwritten Python tests and extensions for Points2D (…

…#2410)

Implements all the necessary Python extensions so that the UX of the
`Points2D` archetype is --dare I say-- _exquisite_.

Also adds a whole bunch of tests for all of these extensions because,
well.. Python...

Might be easier to review commit by commit 🤷 

---
 
Codegen/IDL PR series:
- #2362
- #2363
- #2369
- #2370 
- #2374 
- #2375 
- #2410
- #2432

---------

Co-authored-by: Andreas Reich <[email protected]>

crates/re_types/build.rs

@@ -1,5 +1,7 @@
 //! Generates Rust & Python code from flatbuffers definitions.
 
+use std::path::PathBuf;
+
 use xshell::{cmd, Shell};
 
 use re_build_tools::{
@@ -84,6 +86,15 @@ fn main() {
         "./definitions/rerun/archetypes.fbs",
     );
 
+    let pyproject_path = PathBuf::from(PYTHON_OUTPUT_DIR_PATH)
+        .parent()
+        .unwrap()
+        .parent()
+        .unwrap()
+        .join("pyproject.toml")
+        .to_string_lossy()
+        .to_string();
+
     // NOTE: This requires both `black` and `ruff` to be in $PATH, but only for contributors,
     // not end users.
     // Even for contributors, `black` and `ruff` won't be needed unless they edit some of the
@@ -96,15 +107,25 @@ fn main() {
     // the build.
     //
     // The CI will catch the unformatted files at PR time and complain appropriately anyhow.
-    cmd!(sh, "black {PYTHON_OUTPUT_DIR_PATH}").run().ok();
+    cmd!(
+        sh,
+        "black --config {pyproject_path} {PYTHON_OUTPUT_DIR_PATH}"
+    )
+    .run()
+    .ok();
 
     // NOTE: We're purposefully ignoring the error here.
     //
     // If the user doesn't have `ruff` in their $PATH, there's still no good reason to fail
     // the build.
     //
     // The CI will catch the unformatted files at PR time and complain appropriately anyhow.
-    cmd!(sh, "ruff --fix {PYTHON_OUTPUT_DIR_PATH}").run().ok();
+    cmd!(
+        sh,
+        "ruff --config {pyproject_path} --fix {PYTHON_OUTPUT_DIR_PATH}"
+    )
+    .run()
+    .ok();
 
     write_versioning_hash(SOURCE_HASH_PATH, new_hash);
 }

crates/re_types/definitions/rerun/components/class_id.fbs

@@ -13,8 +13,8 @@ namespace rerun.components;
 /// \rs Used to look up a `crate::components::ClassDescription` within the `crate::components::AnnotationContext`.
 struct ClassId (
   "attr.arrow.transparent",
-  "attr.python.aliases": "float",
-  "attr.python.array_aliases": "npt.NDArray[np.uint8], npt.NDArray[np.uint16], npt.NDArray[np.uint32]",
+  "attr.python.aliases": "int",
+  "attr.python.array_aliases": "npt.NDArray[np.uint8], npt.NDArray[np.uint16], npt.NDArray[np.uint32], npt.NDArray[np.uint64]",
   "attr.rust.derive": "Copy, Clone, Debug, Default, PartialEq, Eq, PartialOrd, Ord, Hash",
   "attr.rust.tuple_struct",
   order: 100

crates/re_types/definitions/rerun/components/color.fbs

@@ -10,10 +10,14 @@ namespace rerun.components;
 // ---
 
 /// An RGBA color tuple with unmultiplied/separate alpha, in sRGB gamma space with linear alpha.
+///
+/// \py Float colors are assumed to be in 0-1 gamma sRGB space.
+/// \py All other colors are assumed to be in 0-255 gamma sRGB space.
+/// \py If there is an alpha, we assume it is in linear space, and separate (NOT pre-multiplied).
 struct Color (
   "attr.arrow.transparent",
-  "attr.python.aliases": "Sequence[int], Sequence[float], npt.NDArray[np.uint8], npt.NDArray[np.float32], npt.NDArray[np.float64]",
-  "attr.python.array_aliases": "Sequence[int], Sequence[float], npt.NDArray[np.uint8], npt.NDArray[np.float32], npt.NDArray[np.float64]",
+  "attr.python.aliases": "int, npt.NDArray[np.uint8], npt.NDArray[np.uint32], npt.NDArray[np.float32], npt.NDArray[np.float64]",
+  "attr.python.array_aliases": "Sequence[int], npt.NDArray[np.uint8], npt.NDArray[np.uint32], npt.NDArray[np.float32], npt.NDArray[np.float64]",
   "attr.rust.derive": "Debug, Default, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, bytemuck::Pod, bytemuck::Zeroable",
   "attr.rust.repr": "transparent",
   "attr.rust.tuple_struct",

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.
-2d611d0d686a6314da94b1797472e7690413465d16a611aae3560a847668d53d
+184f79acfe7c0608f9ec2bb5f6a6783dc2270f5cae0edb147ad08ad4ee89c825

crates/re_types_builder/src/codegen/python.rs

@@ -466,7 +466,7 @@ fn quote_str_repr_from_obj(obj: &Object) -> String {
 
     unindent::unindent(
         r#"
-        def __str__(self):
+        def __str__(self) -> str:
             s = f"rr.{type(self).__name__}(\n"
 
             from dataclasses import fields
@@ -482,7 +482,7 @@ fn quote_str_repr_from_obj(obj: &Object) -> String {
 
             return s
 
-        def __repr__(self):
+        def __repr__(self) -> str:
             return str(self)
 
         "#,
@@ -510,7 +510,7 @@ fn quote_array_method_from_obj(objects: &Objects, obj: &Object) -> String {
     let field_name = &obj.fields[0].name;
     unindent::unindent(&format!(
         "
-        def __array__(self):
+        def __array__(self) -> npt.ArrayLike:
             return np.asarray(self.{field_name})
         ",
     ))
@@ -535,7 +535,7 @@ fn quote_str_method_from_obj(objects: &Objects, obj: &Object) -> String {
     let field_name = &obj.fields[0].name;
     unindent::unindent(&format!(
         "
-        def __str__(self):
+        def __str__(self) -> str:
             return self.{field_name}
         ",
     ))
@@ -771,7 +771,7 @@ fn quote_arrow_support_from_obj(arrow_registry: &ArrowRegistry, obj: &Object) ->
 
                 from .{pkg}_ext import {many}Ext # noqa: E402
 
-                class {arrow}(pa.ExtensionType):
+                class {arrow}(pa.ExtensionType): # type: ignore[misc]
                     def __init__(self: type[pa.ExtensionType]) -> None:
                         pa.ExtensionType.__init__(
                             self, {datatype}, "{fqname}"
@@ -795,7 +795,7 @@ fn quote_arrow_support_from_obj(arrow_registry: &ArrowRegistry, obj: &Object) ->
 
                 class {many}(pa.ExtensionArray, {many}Ext):  # type: ignore[misc]
                     @staticmethod
-                    def from_similar(data: {many_aliases} | None):
+                    def from_similar(data: {many_aliases} | None) -> pa.Array:
                         if data is None:
                             return {arrow}().wrap_array(pa.array([], type={arrow}().storage_type))
                         else:

justfile

@@ -82,7 +82,7 @@ py-lint:
 
 # Run fast unittests
 py-test:
-    python -m pytest rerun_py/tests/unit/
+    python -m pytest -vv rerun_py/tests/unit/
 
 # Serve the python docs locally
 py-docs-serve:

rerun_py/rerun_sdk/rerun2/archetypes/points2d.py

@@ -69,7 +69,7 @@ class Points2D:
     Unique identifiers for each individual point in the batch.
     """
 
-    def __str__(self):
+    def __str__(self) -> str:
         s = f"rr.{type(self).__name__}(\n"
 
         from dataclasses import fields
@@ -86,7 +86,7 @@ def __str__(self):
 
         return s
 
-    def __repr__(self):
+    def __repr__(self) -> str:
         return str(self)
 
     def __init__(

rerun_py/rerun_sdk/rerun2/components/class_id.py

@@ -18,14 +18,19 @@ class ClassId:
 
     id: int
 
-    def __array__(self):
+    def __array__(self) -> npt.ArrayLike:
         return np.asarray(self.id)
 
 
-ClassIdLike = Union[ClassId, float]
+ClassIdLike = Union[ClassId, int]
 
 ClassIdArrayLike = Union[
-    ClassIdLike, Sequence[ClassIdLike], npt.NDArray[np.uint8], npt.NDArray[np.uint16], npt.NDArray[np.uint32]
+    ClassIdLike,
+    Sequence[ClassIdLike],
+    npt.NDArray[np.uint8],
+    npt.NDArray[np.uint16],
+    npt.NDArray[np.uint32],
+    npt.NDArray[np.uint64],
 ]
 
 
@@ -34,7 +39,7 @@ def __array__(self):
 from rerun2.components.class_id_ext import ClassIdArrayExt  # noqa: E402
 
 
-class ClassIdType(pa.ExtensionType):
+class ClassIdType(pa.ExtensionType):  # type: ignore[misc]
     def __init__(self: type[pa.ExtensionType]) -> None:
         pa.ExtensionType.__init__(self, pa.uint16(), "rerun.components.ClassId")
 
@@ -58,7 +63,7 @@ def __arrow_ext_class__(self: type[pa.ExtensionType]) -> type[pa.ExtensionArray]
 
 class ClassIdArray(pa.ExtensionArray, ClassIdArrayExt):  # type: ignore[misc]
     @staticmethod
-    def from_similar(data: ClassIdArrayLike | None):
+    def from_similar(data: ClassIdArrayLike | None) -> pa.Array:
         if data is None:
             return ClassIdType().wrap_array(pa.array([], type=ClassIdType().storage_type))
         else:

rerun_py/rerun_sdk/rerun2/components/class_id_ext.py

@@ -0,0 +1,21 @@
+from __future__ import annotations
+
+__all__ = ["ClassIdArrayExt"]
+
+from typing import Any, Sequence
+
+import numpy as np
+import pyarrow as pa
+
+
+class ClassIdArrayExt:
+    @staticmethod
+    def _from_similar(
+        data: Any | None, *, mono: type, mono_aliases: Any, many: type, many_aliases: Any, arrow: type
+    ) -> pa.Array:
+        if isinstance(data, Sequence) and (len(data) > 0 and isinstance(data[0], mono)):
+            array = np.asarray([class_id.id for class_id in data], np.uint16)
+        else:
+            array = np.asarray(data, dtype=np.uint16).flatten()
+
+        return arrow().wrap_array(pa.array(array, type=arrow().storage_type))

rerun_py/rerun_sdk/rerun2/components/color.py

@@ -14,24 +14,30 @@
 
 @dataclass
 class Color:
-    """An RGBA color tuple with unmultiplied/separate alpha, in sRGB gamma space with linear alpha."""
+    """
+    An RGBA color tuple with unmultiplied/separate alpha, in sRGB gamma space with linear alpha.
+
+    Float colors are assumed to be in 0-1 gamma sRGB space.
+    All other colors are assumed to be in 0-255 gamma sRGB space.
+    If there is an alpha, we assume it is in linear space, and separate (NOT pre-multiplied).
+    """
 
     rgba: int
 
-    def __array__(self):
+    def __array__(self) -> npt.ArrayLike:
         return np.asarray(self.rgba)
 
 
 ColorLike = Union[
-    Color, Sequence[int], Sequence[float], npt.NDArray[np.uint8], npt.NDArray[np.float32], npt.NDArray[np.float64]
+    Color, int, npt.NDArray[np.uint8], npt.NDArray[np.uint32], npt.NDArray[np.float32], npt.NDArray[np.float64]
 ]
 
 ColorArrayLike = Union[
     ColorLike,
     Sequence[ColorLike],
     Sequence[int],
-    Sequence[float],
     npt.NDArray[np.uint8],
+    npt.NDArray[np.uint32],
     npt.NDArray[np.float32],
     npt.NDArray[np.float64],
 ]
@@ -42,7 +48,7 @@ def __array__(self):
 from rerun2.components.color_ext import ColorArrayExt  # noqa: E402
 
 
-class ColorType(pa.ExtensionType):
+class ColorType(pa.ExtensionType):  # type: ignore[misc]
     def __init__(self: type[pa.ExtensionType]) -> None:
         pa.ExtensionType.__init__(self, pa.uint32(), "rerun.components.Color")
 
@@ -66,7 +72,7 @@ def __arrow_ext_class__(self: type[pa.ExtensionType]) -> type[pa.ExtensionArray]
 
 class ColorArray(pa.ExtensionArray, ColorArrayExt):  # type: ignore[misc]
     @staticmethod
-    def from_similar(data: ColorArrayLike | None):
+    def from_similar(data: ColorArrayLike | None) -> pa.Array:
         if data is None:
             return ColorType().wrap_array(pa.array([], type=ColorType().storage_type))
         else:

rerun_py/rerun_sdk/rerun2/components/color_ext.py

@@ -0,0 +1,42 @@
+from __future__ import annotations
+
+__all__ = ["ColorArrayExt"]
+
+from typing import Any, Sequence
+
+import numpy as np
+import pyarrow as pa
+from rerun.color_conversion import u8_array_to_rgba
+
+
+class ColorArrayExt:
+    @staticmethod
+    def _from_similar(
+        data: Any | None, *, mono: type, mono_aliases: Any, many: type, many_aliases: Any, arrow: type
+    ) -> pa.Array:
+        """
+        Normalize flexible colors arrays.
+
+        Float colors are assumed to be in 0-1 gamma sRGB space.
+        All other colors are assumed to be in 0-255 gamma sRGB space.
+
+        If there is an alpha, we assume it is in linear space, and separate (NOT pre-multiplied).
+        """
+        if isinstance(data, Sequence) and len(data) == 0:
+            array = np.array([], np.uint32)
+        elif isinstance(data, Sequence) and (len(data) > 0 and isinstance(data[0], mono)):
+            array = np.asarray([color.rgba for color in data], np.uint32)
+        elif isinstance(data, Sequence) and (len(data) > 0 and isinstance(data[0], int)):
+            array = np.asarray(data, np.uint32)
+        else:
+            array = np.asarray(data)
+            # Rust expects colors in 0-255 uint8
+            if array.dtype.type in [np.float32, np.float64]:
+                # Assume gamma-space colors
+                array = u8_array_to_rgba(np.asarray(np.round(np.asarray(data).reshape((-1, 4)) * 255.0), np.uint8))
+            elif array.dtype.type == np.uint32:
+                array = np.asarray(data).flatten()
+            else:
+                array = u8_array_to_rgba(np.asarray(data, dtype=np.uint8).reshape((-1, 4)))
+
+        return arrow().wrap_array(pa.array(array, type=arrow().storage_type))

rerun_py/rerun_sdk/rerun2/components/draw_order.py

@@ -26,7 +26,7 @@ class DrawOrder:
 
     value: float
 
-    def __array__(self):
+    def __array__(self) -> npt.ArrayLike:
         return np.asarray(self.value)
 
 
@@ -40,7 +40,7 @@ def __array__(self):
 from rerun2.components.draw_order_ext import DrawOrderArrayExt  # noqa: E402
 
 
-class DrawOrderType(pa.ExtensionType):
+class DrawOrderType(pa.ExtensionType):  # type: ignore[misc]
     def __init__(self: type[pa.ExtensionType]) -> None:
         pa.ExtensionType.__init__(self, pa.float32(), "rerun.components.DrawOrder")
 
@@ -64,7 +64,7 @@ def __arrow_ext_class__(self: type[pa.ExtensionType]) -> type[pa.ExtensionArray]
 
 class DrawOrderArray(pa.ExtensionArray, DrawOrderArrayExt):  # type: ignore[misc]
     @staticmethod
-    def from_similar(data: DrawOrderArrayLike | None):
+    def from_similar(data: DrawOrderArrayLike | None) -> pa.Array:
         if data is None:
             return DrawOrderType().wrap_array(pa.array([], type=DrawOrderType().storage_type))
         else:

rerun_py/rerun_sdk/rerun2/components/draw_order_ext.py

@@ -0,0 +1,21 @@
+from __future__ import annotations
+
+__all__ = ["DrawOrderArrayExt"]
+
+from typing import Any, Sequence
+
+import numpy as np
+import pyarrow as pa
+
+
+class DrawOrderArrayExt:
+    @staticmethod
+    def _from_similar(
+        data: Any | None, *, mono: type, mono_aliases: Any, many: type, many_aliases: Any, arrow: type
+    ) -> pa.Array:
+        if isinstance(data, Sequence) and (len(data) > 0 and isinstance(data[0], mono)):
+            array = np.asarray([draw_order.value for draw_order in data], np.float32)
+        else:
+            array = np.require(np.asarray(data), np.float32).flatten()
+
+        return arrow().wrap_array(pa.array(array, type=arrow().storage_type))