feat(encode-types): encode coco type's python classes and test roundtrip (#1030)

Author: georgeh0

python/cocoindex/convert.py

@@ -22,6 +22,7 @@
     AnalyzedTypeInfo,
     AnalyzedUnionType,
     AnalyzedUnknownType,
+    EnrichedValueType,
     analyze_type_info,
     encode_enriched_type,
     is_namedtuple_type,
@@ -611,6 +612,10 @@ def dump_engine_object(v: Any) -> Any:
     """Recursively dump an object for engine. Engine side uses `Pythonized` to catch."""
     if v is None:
         return None
+    elif isinstance(v, EnrichedValueType):
+        return v.encode()
+    elif isinstance(v, FieldSchema):
+        return v.encode()
     elif isinstance(v, type) or get_origin(v) is not None:
         return encode_enriched_type(v)
     elif isinstance(v, Enum):
@@ -660,6 +665,11 @@ def load_engine_object(expected_type: Any, v: Any) -> Any:
     type_info = analyze_type_info(expected_type)
     variant = type_info.variant
 
+    if type_info.core_type is EnrichedValueType:
+        return EnrichedValueType.decode(v)
+    if type_info.core_type is FieldSchema:
+        return FieldSchema.decode(v)
+
     # Any or unknown → return as-is
     if isinstance(variant, AnalyzedAnyType) or type_info.base_type is Any:
         return v

python/cocoindex/tests/test_typing.py

@@ -5,7 +5,6 @@
 from typing import Annotated, Any, Literal, NamedTuple, get_args, get_origin
 
 import numpy as np
-import pytest
 from numpy.typing import NDArray
 
 from cocoindex.typing import (
@@ -20,7 +19,10 @@
     Vector,
     VectorInfo,
     analyze_type_info,
+    decode_engine_value_type,
     encode_enriched_type,
+    encode_enriched_type_info,
+    encode_engine_value_type,
 )
 
 
@@ -32,7 +34,7 @@ class SimpleDataclass:
 
 class SimpleNamedTuple(NamedTuple):
     name: str
-    value: Any
+    value: int
 
 
 def test_ndarray_float32_no_dim() -> None:
@@ -427,3 +429,125 @@ def test_unknown_type() -> None:
     typ = set
     result = analyze_type_info(typ)
     assert isinstance(result.variant, AnalyzedUnknownType)
+
+
+# ========================= Encode/Decode Tests =========================
+
+
+def encode_type_from_annotation(t: Any) -> dict[str, Any]:
+    """Helper function to encode a Python type annotation to its dictionary representation."""
+    return encode_enriched_type_info(analyze_type_info(t))
+
+
+def test_basic_types_encode_decode() -> None:
+    """Test encode/decode roundtrip for basic Python types."""
+    test_cases = [
+        str,
+        int,
+        float,
+        bool,
+        bytes,
+        uuid.UUID,
+        datetime.date,
+        datetime.time,
+        datetime.datetime,
+        datetime.timedelta,
+    ]
+
+    for typ in test_cases:
+        encoded = encode_type_from_annotation(typ)
+        decoded = decode_engine_value_type(encoded["type"])
+        reencoded = encode_engine_value_type(decoded)
+        assert reencoded == encoded["type"]
+
+
+def test_vector_types_encode_decode() -> None:
+    """Test encode/decode roundtrip for vector types."""
+    test_cases = [
+        NDArray[np.float32],
+        NDArray[np.float64],
+        NDArray[np.int64],
+        Vector[np.float32],
+        Vector[np.float32, Literal[128]],
+        Vector[str],
+    ]
+
+    for typ in test_cases:
+        encoded = encode_type_from_annotation(typ)
+        decoded = decode_engine_value_type(encoded["type"])
+        reencoded = encode_engine_value_type(decoded)
+        assert reencoded == encoded["type"]
+
+
+def test_struct_types_encode_decode() -> None:
+    """Test encode/decode roundtrip for struct types."""
+    test_cases = [
+        SimpleDataclass,
+        SimpleNamedTuple,
+    ]
+
+    for typ in test_cases:
+        encoded = encode_type_from_annotation(typ)
+        decoded = decode_engine_value_type(encoded["type"])
+        reencoded = encode_engine_value_type(decoded)
+        assert reencoded == encoded["type"]
+
+
+def test_table_types_encode_decode() -> None:
+    """Test encode/decode roundtrip for table types."""
+    test_cases = [
+        list[SimpleDataclass],  # LTable
+        dict[str, SimpleDataclass],  # KTable
+    ]
+
+    for typ in test_cases:
+        encoded = encode_type_from_annotation(typ)
+        decoded = decode_engine_value_type(encoded["type"])
+        reencoded = encode_engine_value_type(decoded)
+        assert reencoded == encoded["type"]
+
+
+def test_nullable_types_encode_decode() -> None:
+    """Test encode/decode roundtrip for nullable types."""
+    test_cases = [
+        str | None,
+        int | None,
+        NDArray[np.float32] | None,
+    ]
+
+    for typ in test_cases:
+        encoded = encode_type_from_annotation(typ)
+        decoded = decode_engine_value_type(encoded["type"])
+        reencoded = encode_engine_value_type(decoded)
+        assert reencoded == encoded["type"]
+
+
+def test_annotated_types_encode_decode() -> None:
+    """Test encode/decode roundtrip for annotated types."""
+    test_cases = [
+        Annotated[str, TypeAttr("key", "value")],
+        Annotated[NDArray[np.float32], VectorInfo(dim=256)],
+        Annotated[list[int], VectorInfo(dim=10)],
+    ]
+
+    for typ in test_cases:
+        encoded = encode_type_from_annotation(typ)
+        decoded = decode_engine_value_type(encoded["type"])
+        reencoded = encode_engine_value_type(decoded)
+        assert reencoded == encoded["type"]
+
+
+def test_complex_nested_encode_decode() -> None:
+    """Test complex nested structure encode/decode roundtrip."""
+
+    # Create a complex nested structure using Python type annotations
+    @dataclasses.dataclass
+    class ComplexStruct:
+        embedding: NDArray[np.float32]
+        metadata: str | None
+        score: Annotated[float, TypeAttr("indexed", True)]
+
+    encoded = encode_type_from_annotation(ComplexStruct)
+    decoded = decode_engine_value_type(encoded["type"])
+    reencoded = encode_engine_value_type(decoded)
+    assert reencoded == encoded["type"]

python/cocoindex/typing.py

@@ -489,6 +489,12 @@ def decode(obj: dict[str, Any]) -> "VectorTypeSchema":
             dimension=obj.get("dimension"),
         )
 
+    def encode(self) -> dict[str, Any]:
+        return {
+            "element_type": self.element_type.encode(),
+            "dimension": self.dimension,
+        }
+
 
 @dataclasses.dataclass
 class UnionTypeSchema:
@@ -500,6 +506,9 @@ def decode(obj: dict[str, Any]) -> "UnionTypeSchema":
             variants=[BasicValueType.decode(t) for t in obj["types"]]
         )
 
+    def encode(self) -> dict[str, Any]:
+        return {"types": [variant.encode() for variant in self.variants]}
+
 
 @dataclasses.dataclass
 class BasicValueType:
@@ -545,6 +554,14 @@ def decode(obj: dict[str, Any]) -> "BasicValueType":
             )
         return BasicValueType(kind=kind)  # type: ignore[arg-type]
 
+    def encode(self) -> dict[str, Any]:
+        result = {"kind": self.kind}
+        if self.kind == "Vector" and self.vector is not None:
+            result.update(self.vector.encode())
+        elif self.kind == "Union" and self.union is not None:
+            result.update(self.union.encode())
+        return result
+
 
 @dataclasses.dataclass
 class EnrichedValueType:
@@ -560,6 +577,14 @@ def decode(obj: dict[str, Any]) -> "EnrichedValueType":
             attrs=obj.get("attrs"),
         )
 
+    def encode(self) -> dict[str, Any]:
+        result: dict[str, Any] = {"type": self.type.encode()}
+        if self.nullable:
+            result["nullable"] = True
+        if self.attrs is not None:
+            result["attrs"] = self.attrs
+        return result
+
 
 @dataclasses.dataclass
 class FieldSchema:
@@ -570,6 +595,11 @@ class FieldSchema:
     def decode(obj: dict[str, Any]) -> "FieldSchema":
         return FieldSchema(name=obj["name"], value_type=EnrichedValueType.decode(obj))
 
+    def encode(self) -> dict[str, Any]:
+        result = self.value_type.encode()
+        result["name"] = self.name
+        return result
+
 
 @dataclasses.dataclass
 class StructSchema:
@@ -583,11 +613,22 @@ def decode(cls, obj: dict[str, Any]) -> Self:
             description=obj.get("description"),
         )
 
+    def encode(self) -> dict[str, Any]:
+        result: dict[str, Any] = {"fields": [field.encode() for field in self.fields]}
+        if self.description is not None:
+            result["description"] = self.description
+        return result
+
 
 @dataclasses.dataclass
 class StructType(StructSchema):
     kind: Literal["Struct"] = "Struct"
 
+    def encode(self) -> dict[str, Any]:
+        result = super().encode()
+        result["kind"] = self.kind
+        return result
+
 
 @dataclasses.dataclass
 class TableType:
@@ -608,6 +649,12 @@ def decode(obj: dict[str, Any]) -> "TableType":
             num_key_parts=obj.get("num_key_parts"),
         )
 
+    def encode(self) -> dict[str, Any]:
+        result: dict[str, Any] = {"kind": self.kind, "row": self.row.encode()}
+        if self.num_key_parts is not None:
+            result["num_key_parts"] = self.num_key_parts
+        return result
+
 
 ValueType = BasicValueType | StructType | TableType
 
@@ -626,3 +673,8 @@ def decode_engine_value_type(obj: dict[str, Any]) -> ValueType:
 
     # Otherwise it's a basic value
     return BasicValueType.decode(obj)
+
+
+def encode_engine_value_type(value_type: ValueType) -> dict[str, Any]:
+    """Encode a ValueType to its dictionary representation."""
+    return value_type.encode()