cocoindex 0.3.4__cp311-abi3-manylinux_2_28_x86_64.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