cocoindex 0.1.50__cp312-cp312-macosx_11_0_arm64.whl → 0.1.52__cp312-cp312-macosx_11_0_arm64.whl

cocoindex/tests/test_convert.py

@@ -1,12 +1,14 @@
 import uuid
 import datetime
 from dataclasses import dataclass, make_dataclass
-from typing import NamedTuple, Literal, Any, Callable
+from typing import NamedTuple, Literal, Any, Callable, Union
 import pytest
 import cocoindex
 from cocoindex.typing import (
     encode_enriched_type,
     Vector,
+    Float32,
+    Float64,
 )
 from cocoindex.convert import (
     encode_engine_value,
@@ -32,7 +34,7 @@ class Tag:
 
 @dataclass
 class Basket:
-    items: list
+    items: list[str]
 
 
 @dataclass
@@ -74,39 +76,50 @@ def build_engine_value_decoder(
 
 
 def validate_full_roundtrip(
-    value: Any, output_type: Any, input_type: Any | None = None
+    value: Any,
+    value_type: Any = None,
+    *other_decoded_values: tuple[Any, Any],
 ) -> None:
     """
     Validate the given value doesn't change after encoding, sending to engine (using output_type), receiving back and decoding (using input_type).
 
-    If `input_type` is not specified, uses `output_type` as the target.
+    `other_decoded_values` is a tuple of (value, type) pairs.
+    If provided, also validate the value can be decoded to the other types.
     """
-    from cocoindex import _engine
+    from cocoindex import _engine  # type: ignore
 
     encoded_value = encode_engine_value(value)
-    encoded_output_type = encode_enriched_type(output_type)["type"]
+    value_type = value_type or type(value)
+    encoded_output_type = encode_enriched_type(value_type)["type"]
     value_from_engine = _engine.testutil.seder_roundtrip(
         encoded_value, encoded_output_type
     )
-    decoded_value = build_engine_value_decoder(input_type or output_type, output_type)(
+    decoded_value = build_engine_value_decoder(value_type, value_type)(
         value_from_engine
     )
-    assert decoded_value == value
+    np.testing.assert_array_equal(decoded_value, value)
 
+    if other_decoded_values is not None:
+        for other_value, other_type in other_decoded_values:
+            other_decoded_value = build_engine_value_decoder(other_type, other_type)(
+                value_from_engine
+            )
+            np.testing.assert_array_equal(other_decoded_value, other_value)
 
-def test_encode_engine_value_basic_types():
+
+def test_encode_engine_value_basic_types() -> None:
     assert encode_engine_value(123) == 123
     assert encode_engine_value(3.14) == 3.14
     assert encode_engine_value("hello") == "hello"
     assert encode_engine_value(True) is True
 
 
-def test_encode_engine_value_uuid():
+def test_encode_engine_value_uuid() -> None:
     u = uuid.uuid4()
     assert encode_engine_value(u) == u.bytes
 
 
-def test_encode_engine_value_date_time_types():
+def test_encode_engine_value_date_time_types() -> None:
     d = datetime.date(2024, 1, 1)
     assert encode_engine_value(d) == d
     t = datetime.time(12, 30)
@@ -115,7 +128,7 @@ def test_encode_engine_value_date_time_types():
     assert encode_engine_value(dt) == dt
 
 
-def test_encode_engine_value_struct():
+def test_encode_engine_value_struct() -> None:
     order = Order(order_id="O123", name="mixed nuts", price=25.0)
     assert encode_engine_value(order) == ["O123", "mixed nuts", 25.0, "default_extra"]
 
@@ -128,7 +141,7 @@ def test_encode_engine_value_struct():
     ]
 
 
-def test_encode_engine_value_list_of_structs():
+def test_encode_engine_value_list_of_structs() -> None:
     orders = [Order("O1", "item1", 10.0), Order("O2", "item2", 20.0)]
     assert encode_engine_value(orders) == [
         ["O1", "item1", 10.0, "default_extra"],
@@ -145,12 +158,12 @@ def test_encode_engine_value_list_of_structs():
     ]
 
 
-def test_encode_engine_value_struct_with_list():
+def test_encode_engine_value_struct_with_list() -> None:
     basket = Basket(items=["apple", "banana"])
     assert encode_engine_value(basket) == [["apple", "banana"]]
 
 
-def test_encode_engine_value_nested_struct():
+def test_encode_engine_value_nested_struct() -> None:
     customer = Customer(name="Alice", order=Order("O1", "item1", 10.0))
     assert encode_engine_value(customer) == [
         "Alice",
@@ -168,12 +181,12 @@ def test_encode_engine_value_nested_struct():
     ]
 
 
-def test_encode_engine_value_empty_list():
+def test_encode_engine_value_empty_list() -> None:
     assert encode_engine_value([]) == []
     assert encode_engine_value([[]]) == [[]]
 
 
-def test_encode_engine_value_tuple():
+def test_encode_engine_value_tuple() -> None:
     assert encode_engine_value(()) == []
     assert encode_engine_value((1, 2, 3)) == [1, 2, 3]
     assert encode_engine_value(((1, 2), (3, 4))) == [[1, 2], [3, 4]]
@@ -181,20 +194,23 @@ def test_encode_engine_value_tuple():
     assert encode_engine_value(((),)) == [[]]
 
 
-def test_encode_engine_value_none():
+def test_encode_engine_value_none() -> None:
     assert encode_engine_value(None) is None
 
 
-def test_make_engine_value_decoder_basic_types():
-    for engine_type_in_py, value in [
-        (int, 42),
-        (float, 3.14),
-        (str, "hello"),
-        (bool, True),
-        # (type(None), None),  # Removed unsupported NoneType
-    ]:
-        decoder = build_engine_value_decoder(engine_type_in_py)
-        assert decoder(value) == value
+def test_roundtrip_basic_types() -> None:
+    validate_full_roundtrip(42, int)
+    validate_full_roundtrip(3.25, float, (3.25, Float64))
+    validate_full_roundtrip(3.25, Float64, (3.25, float))
+    validate_full_roundtrip(3.25, Float32)
+    validate_full_roundtrip("hello", str)
+    validate_full_roundtrip(True, bool)
+    validate_full_roundtrip(False, bool)
+    validate_full_roundtrip(datetime.date(2025, 1, 1), datetime.date)
+    validate_full_roundtrip(datetime.datetime.now(), cocoindex.LocalDateTime)
+    validate_full_roundtrip(
+        datetime.datetime.now(datetime.UTC), cocoindex.OffsetDateTime
+    )
 
 
 @pytest.mark.parametrize(
@@ -312,18 +328,18 @@ def test_make_engine_value_decoder_basic_types():
         ),
     ],
 )
-def test_struct_decoder_cases(data_type, engine_val, expected):
+def test_struct_decoder_cases(data_type: Any, engine_val: Any, expected: Any) -> None:
     decoder = build_engine_value_decoder(data_type)
     assert decoder(engine_val) == expected
 
 
-def test_make_engine_value_decoder_collections():
+def test_make_engine_value_decoder_list_of_struct() -> None:
     # List of structs (dataclass)
-    decoder = build_engine_value_decoder(list[Order])
     engine_val = [
         ["O1", "item1", 10.0, "default_extra"],
         ["O2", "item2", 20.0, "default_extra"],
     ]
+    decoder = build_engine_value_decoder(list[Order])
     assert decoder(engine_val) == [
         Order("O1", "item1", 10.0, "default_extra"),
         Order("O2", "item2", 20.0, "default_extra"),
@@ -336,13 +352,15 @@ def test_make_engine_value_decoder_collections():
         OrderNamedTuple("O2", "item2", 20.0, "default_extra"),
     ]
 
+
+def test_make_engine_value_decoder_struct_of_list() -> None:
     # Struct with list field
-    decoder = build_engine_value_decoder(Customer)
     engine_val = [
         "Alice",
         ["O1", "item1", 10.0, "default_extra"],
         [["vip"], ["premium"]],
     ]
+    decoder = build_engine_value_decoder(Customer)
     assert decoder(engine_val) == Customer(
         "Alice",
         Order("O1", "item1", 10.0, "default_extra"),
@@ -357,8 +375,9 @@ def test_make_engine_value_decoder_collections():
         [Tag("vip"), Tag("premium")],
     )
 
+
+def test_make_engine_value_decoder_struct_of_struct() -> None:
     # Struct with struct field
-    decoder = build_engine_value_decoder(NestedStruct)
     engine_val = [
         ["Alice", ["O1", "item1", 10.0, "default_extra"], [["vip"]]],
         [
@@ -367,6 +386,7 @@ def test_make_engine_value_decoder_collections():
         ],
         2,
     ]
+    decoder = build_engine_value_decoder(NestedStruct)
     assert decoder(engine_val) == NestedStruct(
         Customer("Alice", Order("O1", "item1", 10.0, "default_extra"), [Tag("vip")]),
         [
@@ -377,11 +397,13 @@ def test_make_engine_value_decoder_collections():
     )
 
 
-def make_engine_order(fields):
+def make_engine_order(fields: list[tuple[str, type]]) -> type:
     return make_dataclass("EngineOrder", fields)
 
 
-def make_python_order(fields, defaults=None):
+def make_python_order(
+    fields: list[tuple[str, type]], defaults: dict[str, Any] | None = None
+) -> type:
     if defaults is None:
         defaults = {}
     # Move all fields with defaults to the end (Python dataclass requirement)
@@ -455,8 +477,12 @@ def make_python_order(fields, defaults=None):
     ],
 )
 def test_field_position_cases(
-    engine_fields, python_fields, python_defaults, engine_val, expected_python_val
-):
+    engine_fields: list[tuple[str, type]],
+    python_fields: list[tuple[str, type]],
+    python_defaults: dict[str, Any],
+    engine_val: list[Any],
+    expected_python_val: tuple[Any, ...],
+) -> None:
     EngineOrder = make_engine_order(engine_fields)
     PythonOrder = make_python_order(python_fields, python_defaults)
     decoder = build_engine_value_decoder(EngineOrder, PythonOrder)
@@ -513,13 +539,13 @@ def test_roundtrip_ktable_struct_key() -> None:
     validate_full_roundtrip(value_nt, t_nt)
 
 
-IntVectorType = cocoindex.Vector[np.int32, Literal[5]]
+IntVectorType = cocoindex.Vector[np.int64, Literal[5]]
 
 
 def test_vector_as_vector() -> None:
-    value: IntVectorType = [1, 2, 3, 4, 5]
+    value = np.array([1, 2, 3, 4, 5], dtype=np.int64)
     encoded = encode_engine_value(value)
-    assert encoded == [1, 2, 3, 4, 5]
+    assert np.array_equal(encoded, value)
     decoded = build_engine_value_decoder(IntVectorType)(encoded)
     assert np.array_equal(decoded, value)
 
@@ -540,12 +566,17 @@ Float32VectorType = Vector[np.float32, Literal[3]]
 Float64VectorType = Vector[np.float64, Literal[3]]
 Int64VectorType = Vector[np.int64, Literal[3]]
 Int32VectorType = Vector[np.int32, Literal[3]]
+UInt8VectorType = Vector[np.uint8, Literal[3]]
+UInt16VectorType = Vector[np.uint16, Literal[3]]
+UInt32VectorType = Vector[np.uint32, Literal[3]]
+UInt64VectorType = Vector[np.uint64, Literal[3]]
+StrVectorType = Vector[str]
 NDArrayFloat32Type = NDArray[np.float32]
 NDArrayFloat64Type = NDArray[np.float64]
 NDArrayInt64Type = NDArray[np.int64]
 
 
-def test_encode_engine_value_ndarray():
+def test_encode_engine_value_ndarray() -> None:
     """Test encoding NDArray vectors to lists for the Rust engine."""
     vec_f32: Float32VectorType = np.array([1.0, 2.0, 3.0], dtype=np.float32)
     assert np.array_equal(encode_engine_value(vec_f32), [1.0, 2.0, 3.0])
@@ -557,7 +588,7 @@ def test_encode_engine_value_ndarray():
     assert np.array_equal(encode_engine_value(vec_nd_f32), [1.0, 2.0, 3.0])
 
 
-def test_make_engine_value_decoder_ndarray():
+def test_make_engine_value_decoder_ndarray() -> None:
     """Test decoding engine lists to NDArray vectors."""
     decoder_f32 = build_engine_value_decoder(Float32VectorType)
     result_f32 = decoder_f32([1.0, 2.0, 3.0])
@@ -581,16 +612,16 @@ def test_make_engine_value_decoder_ndarray():
     assert np.array_equal(result_nd_f32, np.array([1.0, 2.0, 3.0], dtype=np.float32))
 
 
-def test_roundtrip_ndarray_vector():
+def test_roundtrip_ndarray_vector() -> None:
     """Test roundtrip encoding and decoding of NDArray vectors."""
-    value_f32: Float32VectorType = np.array([1.0, 2.0, 3.0], dtype=np.float32)
+    value_f32 = np.array([1.0, 2.0, 3.0], dtype=np.float32)
     encoded_f32 = encode_engine_value(value_f32)
     np.array_equal(encoded_f32, [1.0, 2.0, 3.0])
     decoded_f32 = build_engine_value_decoder(Float32VectorType)(encoded_f32)
     assert isinstance(decoded_f32, np.ndarray)
     assert decoded_f32.dtype == np.float32
     assert np.array_equal(decoded_f32, value_f32)
-    value_i64: Int64VectorType = np.array([1, 2, 3], dtype=np.int64)
+    value_i64 = np.array([1, 2, 3], dtype=np.int64)
     encoded_i64 = encode_engine_value(value_i64)
     assert np.array_equal(encoded_i64, [1, 2, 3])
     decoded_i64 = build_engine_value_decoder(Int64VectorType)(encoded_i64)
@@ -606,63 +637,23 @@ def test_roundtrip_ndarray_vector():
     assert np.array_equal(decoded_nd_f64, value_nd_f64)
 
 
-def test_uint_support():
-    """Test encoding and decoding of unsigned integer vectors."""
-    value_uint8 = np.array([1, 2, 3, 4], dtype=np.uint8)
-    encoded = encode_engine_value(value_uint8)
-    assert np.array_equal(encoded, [1, 2, 3, 4])
-    decoder = make_engine_value_decoder(
-        [], {"kind": "Vector", "element_type": {"kind": "UInt8"}}, NDArray[np.uint8]
-    )
-    decoded = decoder(encoded)
-    assert np.array_equal(decoded, value_uint8)
-    assert decoded.dtype == np.uint8
-    value_uint16 = np.array([1, 2, 3, 4], dtype=np.uint16)
-    encoded = encode_engine_value(value_uint16)
-    assert np.array_equal(encoded, [1, 2, 3, 4])
-    decoder = make_engine_value_decoder(
-        [], {"kind": "Vector", "element_type": {"kind": "UInt16"}}, NDArray[np.uint16]
-    )
-    decoded = decoder(encoded)
-    assert np.array_equal(decoded, value_uint16)
-    assert decoded.dtype == np.uint16
-    value_uint32 = np.array([1, 2, 3], dtype=np.uint32)
-    encoded = encode_engine_value(value_uint32)
-    assert np.array_equal(encoded, [1, 2, 3])
-    decoder = make_engine_value_decoder(
-        [], {"kind": "Vector", "element_type": {"kind": "UInt32"}}, NDArray[np.uint32]
-    )
-    decoded = decoder(encoded)
-    assert np.array_equal(decoded, value_uint32)
-    assert decoded.dtype == np.uint32
-    value_uint64 = np.array([1, 2, 3], dtype=np.uint64)
-    encoded = encode_engine_value(value_uint64)
-    assert np.array_equal(encoded, [1, 2, 3])
-    decoder = make_engine_value_decoder(
-        [], {"kind": "Vector", "element_type": {"kind": "UInt8"}}, NDArray[np.uint64]
-    )
-    decoded = decoder(encoded)
-    assert np.array_equal(decoded, value_uint64)
-    assert decoded.dtype == np.uint64
-
-
-def test_ndarray_dimension_mismatch():
+def test_ndarray_dimension_mismatch() -> None:
     """Test dimension enforcement for Vector with specified dimension."""
-    value: Float32VectorType = np.array([1.0, 2.0], dtype=np.float32)
+    value = np.array([1.0, 2.0], dtype=np.float32)
     encoded = encode_engine_value(value)
     assert np.array_equal(encoded, [1.0, 2.0])
     with pytest.raises(ValueError, match="Vector dimension mismatch"):
         build_engine_value_decoder(Float32VectorType)(encoded)
 
 
-def test_list_vector_backward_compatibility():
+def test_list_vector_backward_compatibility() -> None:
     """Test that list-based vectors still work for backward compatibility."""
-    value: IntVectorType = [1, 2, 3, 4, 5]
+    value = [1, 2, 3, 4, 5]
     encoded = encode_engine_value(value)
     assert encoded == [1, 2, 3, 4, 5]
     decoded = build_engine_value_decoder(IntVectorType)(encoded)
     assert isinstance(decoded, np.ndarray)
-    assert decoded.dtype == np.int32
+    assert decoded.dtype == np.int64
     assert np.array_equal(decoded, np.array([1, 2, 3, 4, 5], dtype=np.int64))
     value_list: ListIntType = [1, 2, 3, 4, 5]
     encoded = encode_engine_value(value_list)
@@ -671,7 +662,7 @@ def test_list_vector_backward_compatibility():
     assert np.array_equal(decoded, [1, 2, 3, 4, 5])
 
 
-def test_encode_complex_structure_with_ndarray():
+def test_encode_complex_structure_with_ndarray() -> None:
     """Test encoding a complex structure that includes an NDArray."""
 
     @dataclass
@@ -684,17 +675,13 @@ def test_encode_complex_structure_with_ndarray():
         name="test_np", data=np.array([1.0, 0.5], dtype=np.float32), value=100
     )
     encoded = encode_engine_value(original)
-    expected = [
-        "test_np",
-        [1.0, 0.5],
-        100,
-    ]
-    assert encoded[0] == expected[0]
-    assert np.array_equal(encoded[1], expected[1])
-    assert encoded[2] == expected[2]
+
+    assert encoded[0] == original.name
+    assert np.array_equal(encoded[1], original.data)
+    assert encoded[2] == original.value
 
 
-def test_decode_nullable_ndarray_none_or_value_input():
+def test_decode_nullable_ndarray_none_or_value_input() -> None:
     """Test decoding a nullable NDArray with None or value inputs."""
     src_type_dict = {
         "kind": "Vector",
@@ -718,7 +705,7 @@ def test_decode_nullable_ndarray_none_or_value_input():
     )
 
 
-def test_decode_vector_string():
+def test_decode_vector_string() -> None:
     """Test decoding a vector of strings works for Python native list type."""
     src_type_dict = {
         "kind": "Vector",
@@ -729,7 +716,7 @@ def test_decode_vector_string():
     assert decoder(["hello", "world"]) == ["hello", "world"]
 
 
-def test_decode_error_non_nullable_or_non_list_vector():
+def test_decode_error_non_nullable_or_non_list_vector() -> None:
     """Test decoding errors for non-nullable vectors or non-list inputs."""
     src_type_dict = {
         "kind": "Vector",
@@ -743,7 +730,7 @@ def test_decode_error_non_nullable_or_non_list_vector():
         decoder("not a list")
 
 
-def test_dump_vector_type_annotation_with_dim():
+def test_dump_vector_type_annotation_with_dim() -> None:
     """Test dumping a vector type annotation with a specified dimension."""
     expected_dump = {
         "type": {
@@ -755,7 +742,7 @@ def test_dump_vector_type_annotation_with_dim():
     assert dump_engine_object(Float32VectorType) == expected_dump
 
 
-def test_dump_vector_type_annotation_no_dim():
+def test_dump_vector_type_annotation_no_dim() -> None:
     """Test dumping a vector type annotation with no dimension."""
     expected_dump_no_dim = {
         "type": {
@@ -765,3 +752,63 @@ def test_dump_vector_type_annotation_no_dim():
         }
     }
     assert dump_engine_object(Float64VectorTypeNoDim) == expected_dump_no_dim
+
+
+def test_full_roundtrip_vector_numeric_types() -> None:
+    """Test full roundtrip for numeric vector types using NDArray."""
+    value_f32: Vector[np.float32, Literal[3]] = np.array(
+        [1.0, 2.0, 3.0], dtype=np.float32
+    )
+    validate_full_roundtrip(value_f32, Vector[np.float32, Literal[3]])
+    value_f64: Vector[np.float64, Literal[3]] = np.array(
+        [1.0, 2.0, 3.0], dtype=np.float64
+    )
+    validate_full_roundtrip(value_f64, Vector[np.float64, Literal[3]])
+    value_i64: Vector[np.int64, Literal[3]] = np.array([1, 2, 3], dtype=np.int64)
+    validate_full_roundtrip(value_i64, Vector[np.int64, Literal[3]])
+    value_i32: Vector[np.int32, Literal[3]] = np.array([1, 2, 3], dtype=np.int32)
+    with pytest.raises(ValueError, match="type unsupported yet"):
+        validate_full_roundtrip(value_i32, Vector[np.int32, Literal[3]])
+    value_u8: Vector[np.uint8, Literal[3]] = np.array([1, 2, 3], dtype=np.uint8)
+    with pytest.raises(ValueError, match="type unsupported yet"):
+        validate_full_roundtrip(value_u8, Vector[np.uint8, Literal[3]])
+    value_u16: Vector[np.uint16, Literal[3]] = np.array([1, 2, 3], dtype=np.uint16)
+    with pytest.raises(ValueError, match="type unsupported yet"):
+        validate_full_roundtrip(value_u16, Vector[np.uint16, Literal[3]])
+    value_u32: Vector[np.uint32, Literal[3]] = np.array([1, 2, 3], dtype=np.uint32)
+    with pytest.raises(ValueError, match="type unsupported yet"):
+        validate_full_roundtrip(value_u32, Vector[np.uint32, Literal[3]])
+    value_u64: Vector[np.uint64, Literal[3]] = np.array([1, 2, 3], dtype=np.uint64)
+    with pytest.raises(ValueError, match="type unsupported yet"):
+        validate_full_roundtrip(value_u64, Vector[np.uint64, Literal[3]])
+
+
+def test_roundtrip_vector_no_dimension() -> None:
+    """Test full roundtrip for vector types without dimension annotation."""
+    value_f64: Vector[np.float64] = np.array([1.0, 2.0, 3.0], dtype=np.float64)
+    validate_full_roundtrip(value_f64, Vector[np.float64])
+
+
+def test_roundtrip_string_vector() -> None:
+    """Test full roundtrip for string vector using list."""
+    value_str: Vector[str] = ["hello", "world"]
+    validate_full_roundtrip(value_str, Vector[str])
+
+
+def test_roundtrip_empty_vector() -> None:
+    """Test full roundtrip for empty numeric vector."""
+    value_empty: Vector[np.float32] = np.array([], dtype=np.float32)
+    validate_full_roundtrip(value_empty, Vector[np.float32])
+
+
+def test_roundtrip_dimension_mismatch() -> None:
+    """Test that dimension mismatch raises an error during roundtrip."""
+    value_f32: Vector[np.float32, Literal[3]] = np.array([1.0, 2.0], dtype=np.float32)
+    with pytest.raises(ValueError, match="Vector dimension mismatch"):
+        validate_full_roundtrip(value_f32, Vector[np.float32, Literal[3]])
+
+
+def test_roundtrip_list_backward_compatibility() -> None:
+    """Test full roundtrip for list-based vectors for backward compatibility."""
+    value_list: list[int] = [1, 2, 3]
+    validate_full_roundtrip(value_list, list[int])