cocoindex 0.2.16__cp311-abi3-manylinux_2_28_aarch64.whl → 0.2.17__cp311-abi3-manylinux_2_28_aarch64.whl

cocoindex/tests/test_engine_object.py

@@ -0,0 +1,331 @@
+import dataclasses
+import datetime
+from typing import TypedDict, NamedTuple, Literal
+
+import numpy as np
+from numpy.typing import NDArray
+import pytest
+
+from cocoindex.typing import Vector
+from cocoindex.engine_object import dump_engine_object, load_engine_object
+
+# Optional Pydantic support for testing
+try:
+    import pydantic
+
+    PYDANTIC_AVAILABLE = True
+except ImportError:
+    PYDANTIC_AVAILABLE = False
+
+
+@dataclasses.dataclass
+class LocalTargetFieldMapping:
+    source: str
+    target: str | None = None
+
+
+@dataclasses.dataclass
+class LocalNodeFromFields:
+    label: str
+    fields: list[LocalTargetFieldMapping]
+
+
+@dataclasses.dataclass
+class LocalNodes:
+    kind = "Node"
+    label: str
+
+
+@dataclasses.dataclass
+class LocalRelationships:
+    kind = "Relationship"
+    rel_type: str
+    source: LocalNodeFromFields
+    target: LocalNodeFromFields
+
+
+class LocalPoint(NamedTuple):
+    x: int
+    y: int
+
+
+class UserInfo(TypedDict):
+    id: str
+    age: int
+
+
+def test_timedelta_roundtrip_via_dump_load() -> None:
+    td = datetime.timedelta(days=1, hours=2, minutes=3, seconds=4, microseconds=500)
+    dumped = dump_engine_object(td)
+    loaded = load_engine_object(datetime.timedelta, dumped)
+    assert isinstance(loaded, datetime.timedelta)
+    assert loaded == td
+
+
+def test_ndarray_roundtrip_via_dump_load() -> None:
+    value: NDArray[np.float32] = np.array([1.0, 2.0, 3.0], dtype=np.float32)
+    dumped = dump_engine_object(value)
+    assert dumped == [1.0, 2.0, 3.0]
+    loaded = load_engine_object(NDArray[np.float32], dumped)
+    assert isinstance(loaded, np.ndarray)
+    assert loaded.dtype == np.float32
+    assert np.array_equal(loaded, value)
+
+
+def test_nodes_kind_is_carried() -> None:
+    node = LocalNodes(label="User")
+    dumped = dump_engine_object(node)
+    # dumped should include discriminator
+    assert dumped.get("kind") == "Node"
+    # load back
+    loaded = load_engine_object(LocalNodes, dumped)
+    assert isinstance(loaded, LocalNodes)
+    # class-level attribute is preserved
+    assert getattr(loaded, "kind", None) == "Node"
+    assert loaded.label == "User"
+
+
+def test_relationships_union_discriminator() -> None:
+    rel = LocalRelationships(
+        rel_type="LIKES",
+        source=LocalNodeFromFields(
+            label="User", fields=[LocalTargetFieldMapping("id")]
+        ),
+        target=LocalNodeFromFields(
+            label="Item", fields=[LocalTargetFieldMapping("id")]
+        ),
+    )
+    dumped = dump_engine_object(rel)
+    assert dumped.get("kind") == "Relationship"
+    loaded = load_engine_object(LocalNodes | LocalRelationships, dumped)
+    assert isinstance(loaded, LocalRelationships)
+    assert getattr(loaded, "kind", None) == "Relationship"
+    assert loaded.rel_type == "LIKES"
+    assert dataclasses.asdict(loaded.source) == {
+        "label": "User",
+        "fields": [{"source": "id", "target": None}],
+    }
+    assert dataclasses.asdict(loaded.target) == {
+        "label": "Item",
+        "fields": [{"source": "id", "target": None}],
+    }
+
+
+def test_typed_dict_roundtrip_via_dump_load() -> None:
+    user: UserInfo = {"id": "u1", "age": 30}
+    dumped = dump_engine_object(user)
+    assert dumped == {"id": "u1", "age": 30}
+    loaded = load_engine_object(UserInfo, dumped)
+    assert loaded == user
+
+
+def test_namedtuple_roundtrip_via_dump_load() -> None:
+    p = LocalPoint(1, 2)
+    dumped = dump_engine_object(p)
+    assert dumped == {"x": 1, "y": 2}
+    loaded = load_engine_object(LocalPoint, dumped)
+    assert isinstance(loaded, LocalPoint)
+    assert loaded == p
+
+
+def test_dataclass_missing_fields_with_auto_defaults() -> None:
+    """Test that missing fields are automatically assigned safe default values."""
+
+    @dataclasses.dataclass
+    class TestClass:
+        required_field: str
+        optional_field: str | None  # Should get None
+        list_field: list[str]  # Should get []
+        dict_field: dict[str, int]  # Should get {}
+        explicit_default: str = "default"  # Should use explicit default
+
+    # Input missing optional_field, list_field, dict_field (but has explicit_default via class definition)
+    input_data = {"required_field": "test_value"}
+
+    loaded = load_engine_object(TestClass, input_data)
+
+    assert isinstance(loaded, TestClass)
+    assert loaded.required_field == "test_value"
+    assert loaded.optional_field is None  # Auto-default for Optional
+    assert loaded.list_field == []  # Auto-default for list
+    assert loaded.dict_field == {}  # Auto-default for dict
+    assert loaded.explicit_default == "default"  # Explicit default from class
+
+
+def test_namedtuple_missing_fields_with_auto_defaults() -> None:
+    """Test that missing fields in NamedTuple are automatically assigned safe default values."""
+    from typing import NamedTuple
+
+    class TestTuple(NamedTuple):
+        required_field: str
+        optional_field: str | None  # Should get None
+        list_field: list[str]  # Should get []
+        dict_field: dict[str, int]  # Should get {}
+
+    # Input missing optional_field, list_field, dict_field
+    input_data = {"required_field": "test_value"}
+
+    loaded = load_engine_object(TestTuple, input_data)
+
+    assert isinstance(loaded, TestTuple)
+    assert loaded.required_field == "test_value"
+    assert loaded.optional_field is None  # Auto-default for Optional
+    assert loaded.list_field == []  # Auto-default for list
+    assert loaded.dict_field == {}  # Auto-default for dict
+
+
+def test_dataclass_unsupported_type_still_fails() -> None:
+    """Test that fields with unsupported types still cause errors when missing."""
+
+    @dataclasses.dataclass
+    class TestClass:
+        required_field1: str
+        required_field2: int  # No auto-default for int
+
+    # Input missing required_field2 which has no safe auto-default
+    input_data = {"required_field1": "test_value"}
+
+    # Should still raise an error because int has no safe auto-default
+    try:
+        load_engine_object(TestClass, input_data)
+        assert False, "Expected TypeError to be raised"
+    except TypeError:
+        pass  # Expected behavior
+
+
+def test_dump_vector_type_annotation_with_dim() -> None:
+    """Test dumping a vector type annotation with a specified dimension."""
+    expected_dump = {
+        "type": {
+            "kind": "Vector",
+            "element_type": {"kind": "Float32"},
+            "dimension": 3,
+        }
+    }
+    assert dump_engine_object(Vector[np.float32, Literal[3]]) == expected_dump
+
+
+def test_dump_vector_type_annotation_no_dim() -> None:
+    """Test dumping a vector type annotation with no dimension."""
+    expected_dump_no_dim = {
+        "type": {
+            "kind": "Vector",
+            "element_type": {"kind": "Float64"},
+            "dimension": None,
+        }
+    }
+    assert dump_engine_object(Vector[np.float64]) == expected_dump_no_dim
+
+
+@pytest.mark.skipif(not PYDANTIC_AVAILABLE, reason="Pydantic not available")
+def test_pydantic_unsupported_type_still_fails() -> None:
+    """Test that fields with unsupported types still cause errors when missing."""
+
+    class TestPydantic(pydantic.BaseModel):
+        required_field1: str
+        required_field2: int  # No auto-default for int
+        optional_field: str | None
+        list_field: list[str]
+        dict_field: dict[str, int]
+        field_with_default: str = "default_value"
+
+    # Input missing required_field2 which has no safe auto-default
+    input_data = {"required_field1": "test_value"}
+
+    # Should still raise an error because int has no safe auto-default
+    with pytest.raises(pydantic.ValidationError):
+        load_engine_object(TestPydantic, input_data)
+
+    assert load_engine_object(
+        TestPydantic, {"required_field1": "test_value", "required_field2": 1}
+    ) == TestPydantic(
+        required_field1="test_value",
+        required_field2=1,
+        field_with_default="default_value",
+        optional_field=None,
+        list_field=[],
+        dict_field={},
+    )
+
+
+@pytest.mark.skipif(not PYDANTIC_AVAILABLE, reason="Pydantic not available")
+def test_pydantic_field_descriptions() -> None:
+    """Test that Pydantic field descriptions are extracted and included in schema."""
+    from pydantic import BaseModel, Field
+
+    class UserWithDescriptions(BaseModel):
+        """A user model with field descriptions."""
+
+        name: str = Field(description="The user's full name")
+        age: int = Field(description="The user's age in years", ge=0, le=150)
+        email: str = Field(description="The user's email address")
+        is_active: bool = Field(
+            description="Whether the user account is active", default=True
+        )
+
+    # Test that field descriptions are extracted
+    encoded_schema = dump_engine_object(UserWithDescriptions)
+
+    # Check that the schema contains field descriptions
+    assert "fields" in encoded_schema["type"]
+    fields = encoded_schema["type"]["fields"]
+
+    # Find fields by name and check descriptions
+    field_descriptions = {field["name"]: field.get("description") for field in fields}
+
+    assert field_descriptions["name"] == "The user's full name"
+    assert field_descriptions["age"] == "The user's age in years"
+    assert field_descriptions["email"] == "The user's email address"
+    assert field_descriptions["is_active"] == "Whether the user account is active"
+
+
+@pytest.mark.skipif(not PYDANTIC_AVAILABLE, reason="Pydantic not available")
+def test_pydantic_field_descriptions_without_field() -> None:
+    """Test that Pydantic models without field descriptions work correctly."""
+    from pydantic import BaseModel
+
+    class UserWithoutDescriptions(BaseModel):
+        """A user model without field descriptions."""
+
+        name: str
+        age: int
+        email: str
+
+    # Test that the schema works without descriptions
+    encoded_schema = dump_engine_object(UserWithoutDescriptions)
+
+    # Check that the schema contains fields but no descriptions
+    assert "fields" in encoded_schema["type"]
+    fields = encoded_schema["type"]["fields"]
+
+    # Verify no descriptions are present
+    for field in fields:
+        assert "description" not in field or field["description"] is None
+
+
+@pytest.mark.skipif(not PYDANTIC_AVAILABLE, reason="Pydantic not available")
+def test_pydantic_mixed_descriptions() -> None:
+    """Test Pydantic model with some fields having descriptions and others not."""
+    from pydantic import BaseModel, Field
+
+    class MixedDescriptions(BaseModel):
+        """A model with mixed field descriptions."""
+
+        name: str = Field(description="The name field")
+        age: int  # No description
+        email: str = Field(description="The email field")
+        active: bool  # No description
+
+    # Test that only fields with descriptions have them in the schema
+    encoded_schema = dump_engine_object(MixedDescriptions)
+
+    assert "fields" in encoded_schema["type"]
+    fields = encoded_schema["type"]["fields"]
+
+    # Find fields by name and check descriptions
+    field_descriptions = {field["name"]: field.get("description") for field in fields}
+
+    assert field_descriptions["name"] == "The name field"
+    assert field_descriptions["age"] is None
+    assert field_descriptions["email"] == "The email field"
+    assert field_descriptions["active"] is None

cocoindex/tests/{test_convert.py → test_engine_value.py}

@@ -8,9 +8,18 @@ import numpy as np
 import pytest
 from numpy.typing import NDArray
 
+# Optional Pydantic support for testing
+try:
+    from pydantic import BaseModel, Field
+
+    PYDANTIC_AVAILABLE = True
+except ImportError:
+    BaseModel = None  # type: ignore[misc,assignment]
+    Field = None  # type: ignore[misc,assignment]
+    PYDANTIC_AVAILABLE = False
+
 import cocoindex
-from cocoindex.convert import (
-    dump_engine_object,
+from cocoindex.engine_value import (
     make_engine_value_encoder,
     make_engine_value_decoder,
 )
@@ -70,6 +79,29 @@ class CustomerNamedTuple(NamedTuple):
     tags: list[Tag] | None = None
 
 
+# Pydantic model definitions (if available)
+if PYDANTIC_AVAILABLE:
+
+    class OrderPydantic(BaseModel):
+        order_id: str
+        name: str
+        price: float
+        extra_field: str = "default_extra"
+
+    class TagPydantic(BaseModel):
+        name: str
+
+    class CustomerPydantic(BaseModel):
+        name: str
+        order: OrderPydantic
+        tags: list[TagPydantic] | None = None
+
+    class NestedStructPydantic(BaseModel):
+        customer: CustomerPydantic
+        orders: list[OrderPydantic]
+        count: int = 0
+
+
 def encode_engine_value(value: Any, type_hint: Type[Any] | str) -> Any:
     """
     Encode a Python value to an engine value.
@@ -973,30 +1005,6 @@ def test_decode_error_non_nullable_or_non_list_vector() -> None:
         decoder("not a list")
 
 
-def test_dump_vector_type_annotation_with_dim() -> None:
-    """Test dumping a vector type annotation with a specified dimension."""
-    expected_dump = {
-        "type": {
-            "kind": "Vector",
-            "element_type": {"kind": "Float32"},
-            "dimension": 3,
-        }
-    }
-    assert dump_engine_object(Float32VectorType) == expected_dump
-
-
-def test_dump_vector_type_annotation_no_dim() -> None:
-    """Test dumping a vector type annotation with no dimension."""
-    expected_dump_no_dim = {
-        "type": {
-            "kind": "Vector",
-            "element_type": {"kind": "Float64"},
-            "dimension": None,
-        }
-    }
-    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 = np.array([1.0, 2.0, 3.0], dtype=np.float32)
@@ -1566,3 +1574,119 @@ def test_auto_default_for_supported_and_unsupported_types() -> None:
         match=r"Field 'b' \(type <class 'int'>\) without default value is missing in input: ",
     ):
         build_engine_value_decoder(Base, UnsupportedField)
+
+
+# Pydantic model tests
+@pytest.mark.skipif(not PYDANTIC_AVAILABLE, reason="Pydantic not available")
+def test_pydantic_simple_struct() -> None:
+    """Test basic Pydantic model encoding and decoding."""
+    order = OrderPydantic(order_id="O1", name="item1", price=10.0)
+    validate_full_roundtrip(order, OrderPydantic)
+
+
+@pytest.mark.skipif(not PYDANTIC_AVAILABLE, reason="Pydantic not available")
+def test_pydantic_struct_with_defaults() -> None:
+    """Test Pydantic model with default values."""
+    order = OrderPydantic(order_id="O1", name="item1", price=10.0)
+    assert order.extra_field == "default_extra"
+    validate_full_roundtrip(order, OrderPydantic)
+
+
+@pytest.mark.skipif(not PYDANTIC_AVAILABLE, reason="Pydantic not available")
+def test_pydantic_nested_struct() -> None:
+    """Test nested Pydantic models."""
+    order = OrderPydantic(order_id="O1", name="item1", price=10.0)
+    customer = CustomerPydantic(name="Alice", order=order)
+    validate_full_roundtrip(customer, CustomerPydantic)
+
+
+@pytest.mark.skipif(not PYDANTIC_AVAILABLE, reason="Pydantic not available")
+def test_pydantic_struct_with_list() -> None:
+    """Test Pydantic model with list fields."""
+    order = OrderPydantic(order_id="O1", name="item1", price=10.0)
+    tags = [TagPydantic(name="vip"), TagPydantic(name="premium")]
+    customer = CustomerPydantic(name="Alice", order=order, tags=tags)
+    validate_full_roundtrip(customer, CustomerPydantic)
+
+
+@pytest.mark.skipif(not PYDANTIC_AVAILABLE, reason="Pydantic not available")
+def test_pydantic_complex_nested_struct() -> None:
+    """Test complex nested Pydantic structure."""
+    order1 = OrderPydantic(order_id="O1", name="item1", price=10.0)
+    order2 = OrderPydantic(order_id="O2", name="item2", price=20.0)
+    customer = CustomerPydantic(
+        name="Alice", order=order1, tags=[TagPydantic(name="vip")]
+    )
+    nested = NestedStructPydantic(customer=customer, orders=[order1, order2], count=2)
+    validate_full_roundtrip(nested, NestedStructPydantic)
+
+
+@pytest.mark.skipif(not PYDANTIC_AVAILABLE, reason="Pydantic not available")
+def test_pydantic_struct_to_dict_binding() -> None:
+    """Test Pydantic model -> dict binding."""
+    order = OrderPydantic(order_id="O1", name="item1", price=10.0, extra_field="custom")
+    expected_dict = {
+        "order_id": "O1",
+        "name": "item1",
+        "price": 10.0,
+        "extra_field": "custom",
+    }
+
+    validate_full_roundtrip(
+        order,
+        OrderPydantic,
+        (expected_dict, Any),
+        (expected_dict, dict),
+        (expected_dict, dict[Any, Any]),
+        (expected_dict, dict[str, Any]),
+    )
+
+
+@pytest.mark.skipif(not PYDANTIC_AVAILABLE, reason="Pydantic not available")
+def test_make_engine_value_decoder_pydantic_struct() -> None:
+    """Test engine value decoder for Pydantic models."""
+    engine_val = ["O1", "item1", 10.0, "default_extra"]
+    decoder = build_engine_value_decoder(OrderPydantic)
+    result = decoder(engine_val)
+
+    assert isinstance(result, OrderPydantic)
+    assert result.order_id == "O1"
+    assert result.name == "item1"
+    assert result.price == 10.0
+    assert result.extra_field == "default_extra"
+
+
+@pytest.mark.skipif(not PYDANTIC_AVAILABLE, reason="Pydantic not available")
+def test_make_engine_value_decoder_pydantic_nested() -> None:
+    """Test engine value decoder for nested Pydantic models."""
+    engine_val = [
+        "Alice",
+        ["O1", "item1", 10.0, "default_extra"],
+        [["vip"]],
+    ]
+    decoder = build_engine_value_decoder(CustomerPydantic)
+    result = decoder(engine_val)
+
+    assert isinstance(result, CustomerPydantic)
+    assert result.name == "Alice"
+    assert isinstance(result.order, OrderPydantic)
+    assert result.order.order_id == "O1"
+    assert result.tags is not None
+    assert len(result.tags) == 1
+    assert isinstance(result.tags[0], TagPydantic)
+    assert result.tags[0].name == "vip"
+
+
+@pytest.mark.skipif(not PYDANTIC_AVAILABLE, reason="Pydantic not available")
+def test_pydantic_mixed_with_dataclass() -> None:
+    """Test mixing Pydantic models with dataclasses."""
+
+    # Create a dataclass that uses a Pydantic model
+    @dataclass
+    class MixedStruct:
+        name: str
+        pydantic_order: OrderPydantic
+
+    order = OrderPydantic(order_id="O1", name="item1", price=10.0)
+    mixed = MixedStruct(name="test", pydantic_order=order)
+    validate_full_roundtrip(mixed, MixedStruct)