Update union type

Author: chardoncs

python/cocoindex/typing.py

@@ -233,7 +233,7 @@ def _encode_type(type_info: AnalyzedTypeInfo) -> dict[str, Any]:
         encoded_type['dimension'] = type_info.vector_info.dim
 
     elif type_info.kind == 'Union':
-        if type_info.elem_type is not UnionType:
+        if type_info.elem_type is not types.UnionType:
             raise ValueError("Union type must have a union-typed element type")
         encoded_type['element_type'] = [
             _encode_type(analyze_type_info(typ)) for typ in typing.get_args(type_info.elem_type)

src/base/schema.rs

@@ -437,58 +437,57 @@ pub struct OpArgSchema {
 #[cfg(test)]
 mod tests {
     use super::*;
+    use utils::union::UnionParseResult;
 
     #[test]
     fn test_union_fmt_empty() {
-        let typ = BasicValueType::Union(UnionType::default());
-        let expected = "Union[]";
-
-        assert_eq!(typ.to_string(), expected);
+        let result = UnionType::parse_from(Vec::new());
+        assert!(matches!(result, Err(_)));
     }
 
     #[test]
     fn test_union_fmt_single() {
-        let typ = BasicValueType::Union(vec![BasicValueType::Uuid].into());
-        let expected = "Union[Uuid]";
-
-        assert_eq!(typ.to_string(), expected);
+        let result = UnionType::parse_from(vec![BasicValueType::Uuid]).unwrap();
+        assert!(matches!(result, UnionParseResult::Single(BasicValueType::Uuid)));
     }
 
     #[test]
-    fn test_union_fmt_autosort_basic() {
-        // Uuid | Date | Str | Bytes = Bytes | Str | Uuid | Date
-        let typ = BasicValueType::Union(vec![
+    fn test_union_fmt_flat() {
+        let union = UnionType::coerce_from(vec![
             BasicValueType::Uuid,
             BasicValueType::Date,
             BasicValueType::Str,
             BasicValueType::Bytes,
-        ].into());
+        ]);
+
+        // Uuid | Date | Str | Bytes = Bytes | Str | Uuid | Date
+        let typ = BasicValueType::Union(union);
         let expected = "Union[Bytes | Str | Uuid | Date]";
 
         assert_eq!(typ.to_string(), expected);
     }
 
     #[test]
-    fn test_union_fmt_nested_auto_unpack() {
+    fn test_union_fmt_nested() {
         // Uuid | Date | (Date | OffsetDateTime | Time | Bytes | (Bytes | Uuid | Time)) | Str |
         // Bytes = Bytes | Str | Uuid | Date | Time | OffsetDateTime
-        let typ = BasicValueType::Union(vec![
+        let typ = BasicValueType::Union(UnionType::coerce_from(vec![
             BasicValueType::Uuid,
             BasicValueType::Date,
-            BasicValueType::Union(vec![
+            BasicValueType::Union(UnionType::coerce_from(vec![
                 BasicValueType::Date,
                 BasicValueType::OffsetDateTime,
                 BasicValueType::Time,
                 BasicValueType::Bytes,
-                BasicValueType::Union(vec![
+                BasicValueType::Union(UnionType::coerce_from(vec![
                     BasicValueType::Bytes,
                     BasicValueType::Uuid,
                     BasicValueType::Time,
-                ].into()),
-            ].into()),
+                ])),
+            ])),
             BasicValueType::Str,
             BasicValueType::Bytes,
-        ].into());
+        ]));
         let expected = "Union[Bytes | Str | Uuid | Date | Time | OffsetDateTime]";
 
         assert_eq!(typ.to_string(), expected);

src/utils/union.rs

@@ -2,13 +2,23 @@ use std::{str::FromStr, sync::Arc};
 
 use crate::{base::{schema::BasicValueType, value::BasicValue}, prelude::*};
 
+#[derive(Debug, Clone)]
+pub enum UnionParseResult {
+    Union(UnionType),
+    Single(BasicValueType),
+}
+
 /// Union type helper storing an auto-sorted set of types excluding `Union`
-#[derive(Debug, Clone, Default, Serialize, Deserialize, PartialEq, Eq, PartialOrd, Ord)]
+#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq, PartialOrd, Ord)]
 pub struct UnionType {
     types: BTreeSet<BasicValueType>,
 }
 
 impl UnionType {
+    fn new() -> Self {
+        Self { types: BTreeSet::new() }
+    }
+
     pub fn types(&self) -> &BTreeSet<BasicValueType> {
         &self.types
     }
@@ -31,32 +41,56 @@ impl UnionType {
         }
     }
 
-    pub fn unpack(self) -> Self {
-        self.types.into()
-    }
-}
+    fn resolve(self) -> Result<UnionParseResult> {
+        if self.types().is_empty() {
+            anyhow::bail!("The union is empty");
+        }
 
-impl From<Vec<BasicValueType>> for UnionType {
-    fn from(value: Vec<BasicValueType>) -> Self {
-        let mut union = Self::default();
+        if self.types().len() == 1 {
+            let mut type_tree: BTreeSet<BasicValueType> = self.into();
+            return Ok(UnionParseResult::Single(type_tree.pop_first().unwrap()));
+        }
 
-        for typ in value {
+        Ok(UnionParseResult::Union(self))
+    }
+
+    /// Move an iterable and parse it into a union type.
+    /// If there is only one single unique type, it returns a single `BasicValueType`.
+    pub fn parse_from<T>(
+        input: impl IntoIterator<Item = BasicValueType, IntoIter = T>,
+    ) -> Result<UnionParseResult>
+    where
+        T: Iterator<Item = BasicValueType>,
+    {
+        let mut union = Self::new();
+
+        for typ in input {
             union.insert(typ);
         }
 
-        union
+        union.resolve()
     }
-}
-
-impl From<BTreeSet<BasicValueType>> for UnionType {
-    fn from(value: BTreeSet<BasicValueType>) -> Self {
-        let mut union = Self::default();
 
-        for typ in value {
-            union.insert(typ);
+    /// Assume the input already contains multiple unique types, panic otherwise.
+    ///
+    /// This method is meant for streamlining the code for test cases.
+    /// Use `parse_from()` instead unless you know the input.
+    pub fn coerce_from<T>(
+        input: impl IntoIterator<Item = BasicValueType, IntoIter = T>,
+    ) -> Self
+    where
+        T: Iterator<Item = BasicValueType>,
+    {
+        match Self::parse_from(input) {
+            Ok(UnionParseResult::Union(union)) => union,
+            _ => panic!("Do not use `coerce_from()` for basic type lists that can possibly be one type."),
         }
+    }
+}
 
-        union
+impl Into<BTreeSet<BasicValueType>> for UnionType {
+    fn into(self) -> BTreeSet<BasicValueType> {
+        self.types
     }
 }