Python codegen: use | for union types

crates/re_types_builder/src/codegen/python.rs

@@ -789,9 +789,9 @@ fn code_for_union(
     };
 
     let inner_type = if field_types.len() > 1 {
-        format!("Union[{}]", field_types.iter().join(", "))
+        field_types.iter().join(" | ")
     } else {
-        field_types.iter().next().unwrap().to_string()
+        field_types.first().unwrap().to_string()
     };
 
     // components and datatypes have converters only if manually provided
@@ -1060,44 +1060,39 @@ fn quote_native_types_method_from_obj(objects: &Objects, obj: &Object) -> String
 fn quote_aliases_from_object(obj: &Object) -> String {
     assert_ne!(obj.kind, ObjectKind::Archetype);
 
-    let aliases = obj.try_get_attr::<String>(ATTR_PYTHON_ALIASES);
-    let array_aliases = obj
-        .try_get_attr::<String>(ATTR_PYTHON_ARRAY_ALIASES)
-        .unwrap_or_default();
+    let aliases = obj.get_attr_list(ATTR_PYTHON_ALIASES);
+    let array_aliases = obj.get_attr_list(ATTR_PYTHON_ARRAY_ALIASES);
 
     let name = &obj.name;
 
     let mut code = String::new();
 
     code.push_unindented_text(
-        &if let Some(aliases) = aliases {
+        &if aliases.is_empty() {
+            format!("{name}Like = {name}")
+        } else {
             format!(
                 r#"
                 if TYPE_CHECKING:
-                    {name}Like = Union[
-                        {name},
-                        {aliases}
-                    ]
+                    {name}Like = {name} | {}
                 else:
                     {name}Like = Any
                 "#,
+                aliases.iter().join(" | ")
             )
-        } else {
-            format!("{name}Like = {name}")
         },
         1,
     );
 
     code.push_unindented_text(
-        format!(
-            r#"
-            {name}ArrayLike = Union[
-                {name},
-                Sequence[{name}Like],
-                {array_aliases}
-            ]
-            "#,
-        ),
+        if array_aliases.is_empty() {
+            format!("{name}ArrayLike = {name} | Sequence[{name}Like]")
+        } else {
+            format!(
+                "{name}ArrayLike = {name} | Sequence[{name}Like] | {}",
+                array_aliases.iter().join(" | ")
+            )
+        },
         0,
     );
 
@@ -1108,34 +1103,33 @@ fn quote_aliases_from_object(obj: &Object) -> String {
 /// included.
 fn quote_union_aliases_from_object<'a>(
     obj: &Object,
-    mut field_types: impl Iterator<Item = &'a String>,
+    field_types: impl Iterator<Item = &'a String>,
 ) -> String {
     assert_ne!(obj.kind, ObjectKind::Archetype);
 
-    let aliases = obj.try_get_attr::<String>(ATTR_PYTHON_ALIASES);
-    let array_aliases = obj
-        .try_get_attr::<String>(ATTR_PYTHON_ARRAY_ALIASES)
-        .unwrap_or_default();
+    let field_types = field_types.collect_vec();
+
+    let aliases = obj.get_attr_list(ATTR_PYTHON_ALIASES);
+    let array_aliases = obj.get_attr_list(ATTR_PYTHON_ARRAY_ALIASES);
 
     let name = &obj.name;
 
-    let union_fields = field_types.join(",");
-    let aliases = if let Some(aliases) = aliases {
-        aliases
-    } else {
-        String::new()
-    };
+    let like = std::iter::once(name)
+        .chain(field_types.iter().copied())
+        .chain(aliases.iter())
+        .join(" | ");
+
+    let array_like = std::iter::once(name)
+        .chain(field_types)
+        .chain(std::iter::once(&format!("Sequence[{name}Like]")))
+        .chain(array_aliases.iter())
+        .join(" | ");
 
     unindent::unindent(&format!(
         r#"
             if TYPE_CHECKING:
-                {name}Like = Union[
-                    {name},{union_fields},{aliases}
-                ]
-                {name}ArrayLike = Union[
-                    {name},{union_fields},
-                    Sequence[{name}Like],{array_aliases}
-                ]
+                {name}Like = {like}
+                {name}ArrayLike = {array_like}
             else:
                 {name}Like = Any
                 {name}ArrayLike = Any

crates/re_types_builder/src/objects.rs

@@ -612,6 +612,15 @@ impl Object {
         self.attrs.try_get(self.fqname.as_str(), name)
     }
 
+    pub fn get_attr_list(&self, name: &str) -> Vec<String> {
+        split_list(
+            &self
+                .attrs
+                .try_get::<String>(self.fqname.as_str(), name)
+                .unwrap_or_default(),
+        )
+    }
+
     pub fn is_attr_set(&self, name: impl AsRef<str>) -> bool {
         self.attrs.has(name)
     }
@@ -1333,3 +1342,77 @@ fn filepath_from_declaration_file(
     .canonicalize_utf8()
     .expect("Failed to canonicalize declaration path")
 }
+
+fn split_list(str: &str) -> Vec<String> {
+    let mut bow = 0;
+    let mut angles = 0;
+    let mut curly = 0;
+    let mut square = 0;
+
+    let mut results = Vec::new();
+    let mut current = String::new();
+
+    for c in str.chars() {
+        match c {
+            '<' => angles += 1,
+            '>' => angles -= 1,
+
+            '(' => bow += 1,
+            ')' => bow -= 1,
+
+            '{' => curly += 1,
+            '}' => curly -= 1,
+
+            '[' => square += 1,
+            ']' => square -= 1,
+
+            ',' if angles == 0 && bow == 0 && curly == 0 && square == 0 => {
+                results.push(current.trim().to_owned());
+                current = String::new();
+                continue; // skip the comma
+            }
+            _ => {}
+        }
+        current.push(c);
+    }
+
+    current = current.trim().to_owned();
+    if !current.is_empty() {
+        results.push(current);
+    }
+    results
+}
+
+#[test]
+fn test_split_list() {
+    assert_eq!(split_list(""), Vec::<String>::new());
+    assert_eq!(split_list(" hello "), vec!["hello".to_owned()]);
+    assert_eq!(
+        split_list("hello, there"),
+        vec!["hello".to_owned(), "there".to_owned()]
+    );
+
+    // Python:
+    assert_eq!(
+        split_list("int, Tuple[int, str], Tuple[int, str, datatypes.Rgba32Like]"),
+        vec![
+            "int".to_owned(),
+            "Tuple[int, str]".to_owned(),
+            "Tuple[int, str, datatypes.Rgba32Like]".to_owned(),
+        ]
+    );
+    assert_eq!(
+        split_list("int, Sequence[Union[int, float]], npt.NDArray[Union[np.uint8, np.float32, np.float64]]"),
+        vec![
+            "int".to_owned(),
+            "Sequence[Union[int, float]]".to_owned(),
+            "npt.NDArray[Union[np.uint8, np.float32, np.float64]]".to_owned(),
+        ]
+    );
+
+    // Rust:
+    assert_eq!(
+        split_list("HashMap<int, String>, bool"),
+        vec!["HashMap<int, String>".to_owned(), "bool".to_owned()]
+    );
+}