Implement fix in codegen + add test

rerun-io · Oct 11, 2024 · bb3f99f · bb3f99f
1 parent eab458a
commit bb3f99f
Show file tree

Hide file tree

Showing 9 changed files with 65 additions and 34 deletions.
diff --git a/crates/build/re_types_builder/src/codegen/python/mod.rs b/crates/build/re_types_builder/src/codegen/python/mod.rs
@@ -1995,9 +1995,11 @@ fn quote_arrow_serialization(
                     return Ok(unindent(
                         r##"
                             if isinstance(data, str):
-                                array = [data]
+                                array: Union[list[str], npt.ArrayLike] = [data]
                             elif isinstance(data, Sequence):
                                 array = [str(datum) for datum in data]
+                            elif isinstance(data, np.ndarray):
+                                array = data
                             else:
                                 array = [str(data)]
 

diff --git a/crates/store/re_types/definitions/rerun/datatypes/utf8.fbs b/crates/store/re_types/definitions/rerun/datatypes/utf8.fbs
@@ -8,7 +8,7 @@ namespace rerun.datatypes;
 table Utf8 (
   "attr.arrow.transparent",
   "attr.python.aliases": "str",
-  "attr.python.array_aliases": "str, Sequence[str]",
+  "attr.python.array_aliases": "str, Sequence[str], npt.ArrayLike",
   "attr.rust.derive": "Default, PartialEq, Eq, PartialOrd, Ord, Hash",
   "attr.rust.override_crate": "re_types_core",
   "attr.rust.repr": "transparent",

diff --git a/rerun_py/rerun_sdk/rerun/datatypes/entity_path.py b/rerun_py/rerun_sdk/rerun/datatypes/entity_path.py
diff --git a/rerun_py/rerun_sdk/rerun/datatypes/utf8.py b/rerun_py/rerun_sdk/rerun/datatypes/utf8.py
diff --git a/rerun_py/rerun_sdk/rerun/datatypes/utf8_ext.py b/rerun_py/rerun_sdk/rerun/datatypes/utf8_ext.py
diff --git a/rerun_py/tests/test_types/components/affix_fuzzer10.py b/rerun_py/tests/test_types/components/affix_fuzzer10.py
diff --git a/rerun_py/tests/test_types/components/affix_fuzzer9.py b/rerun_py/tests/test_types/components/affix_fuzzer9.py
diff --git a/rerun_py/tests/test_types/datatypes/string_component.py b/rerun_py/tests/test_types/datatypes/string_component.py
diff --git a/rerun_py/tests/unit/test_utf8.py b/rerun_py/tests/unit/test_utf8.py
@@ -0,0 +1,27 @@
+from __future__ import annotations
+
+import numpy as np
+from rerun import datatypes
+
+
+def test_utf8_batch_single() -> None:
+    single_string = "hello"
+    list_of_one_string = ["hello"]
+    array_of_one_string = np.array(["hello"])
+
+    assert (
+        datatypes.Utf8Batch(single_string).as_arrow_array() == datatypes.Utf8Batch(list_of_one_string).as_arrow_array()
+    )
+
+    assert (
+        datatypes.Utf8Batch(single_string).as_arrow_array() == datatypes.Utf8Batch(array_of_one_string).as_arrow_array()
+    )
+
+
+def test_utf8_batch_many() -> None:
+    list_of_strings = ["hello", "world"]
+    array_of_strings = np.array(["hello", "world"])
+
+    assert (
+        datatypes.Utf8Batch(list_of_strings).as_arrow_array() == datatypes.Utf8Batch(array_of_strings).as_arrow_array()
+    )