cocoindex 0.1.49__cp311-cp311-win_amd64.whl → 0.1.51__cp311-cp311-win_amd64.whl

cocoindex/tests/test_typing.py

@@ -0,0 +1,505 @@
+import dataclasses
+import datetime
+import uuid
+from typing import (
+    Annotated,
+    List,
+    Dict,
+    Literal,
+    Any,
+    get_args,
+    NamedTuple,
+)
+from collections.abc import Sequence, Mapping
+import pytest
+import numpy as np
+from numpy.typing import NDArray
+
+from cocoindex.typing import (
+    analyze_type_info,
+    Vector,
+    VectorInfo,
+    TypeKind,
+    TypeAttr,
+    Float32,
+    Float64,
+    encode_enriched_type,
+    AnalyzedTypeInfo,
+)
+
+
+@dataclasses.dataclass
+class SimpleDataclass:
+    name: str
+    value: int
+
+
+class SimpleNamedTuple(NamedTuple):
+    name: str
+    value: Any
+
+
+def test_ndarray_float32_no_dim() -> None:
+    typ = NDArray[np.float32]
+    result = analyze_type_info(typ)
+    assert result == AnalyzedTypeInfo(
+        kind="Vector",
+        vector_info=VectorInfo(dim=None),
+        elem_type=Float32,
+        key_type=None,
+        struct_type=None,
+        np_number_type=np.float32,
+        attrs=None,
+        nullable=False,
+    )
+
+
+def test_vector_float32_no_dim() -> None:
+    typ = Vector[np.float32]
+    result = analyze_type_info(typ)
+    assert result == AnalyzedTypeInfo(
+        kind="Vector",
+        vector_info=VectorInfo(dim=None),
+        elem_type=Float32,
+        key_type=None,
+        struct_type=None,
+        np_number_type=np.float32,
+        attrs=None,
+        nullable=False,
+    )
+
+
+def test_ndarray_float64_with_dim() -> None:
+    typ = Annotated[NDArray[np.float64], VectorInfo(dim=128)]
+    result = analyze_type_info(typ)
+    assert result == AnalyzedTypeInfo(
+        kind="Vector",
+        vector_info=VectorInfo(dim=128),
+        elem_type=Float64,
+        key_type=None,
+        struct_type=None,
+        np_number_type=np.float64,
+        attrs=None,
+        nullable=False,
+    )
+
+
+def test_vector_float32_with_dim() -> None:
+    typ = Vector[np.float32, Literal[384]]
+    result = analyze_type_info(typ)
+    assert result == AnalyzedTypeInfo(
+        kind="Vector",
+        vector_info=VectorInfo(dim=384),
+        elem_type=Float32,
+        key_type=None,
+        struct_type=None,
+        np_number_type=np.float32,
+        attrs=None,
+        nullable=False,
+    )
+
+
+def test_ndarray_int64_no_dim() -> None:
+    typ = NDArray[np.int64]
+    result = analyze_type_info(typ)
+    assert result.kind == "Vector"
+    assert result.vector_info == VectorInfo(dim=None)
+    assert get_args(result.elem_type) == (int, TypeKind("Int64"))
+    assert not result.nullable
+
+
+def test_nullable_ndarray() -> None:
+    typ = NDArray[np.float32] | None
+    result = analyze_type_info(typ)
+    assert result == AnalyzedTypeInfo(
+        kind="Vector",
+        vector_info=VectorInfo(dim=None),
+        elem_type=Float32,
+        key_type=None,
+        struct_type=None,
+        np_number_type=np.float32,
+        attrs=None,
+        nullable=True,
+    )
+
+
+def test_vector_str() -> None:
+    typ = Vector[str]
+    result = analyze_type_info(typ)
+    assert result.kind == "Vector"
+    assert result.elem_type == str
+    assert result.vector_info == VectorInfo(dim=None)
+
+
+def test_vector_complex64() -> None:
+    typ = Vector[np.complex64]
+    result = analyze_type_info(typ)
+    assert result.kind == "Vector"
+    assert result.elem_type == np.complex64
+    assert result.vector_info == VectorInfo(dim=None)
+
+
+def test_non_numpy_vector() -> None:
+    typ = Vector[float, Literal[3]]
+    result = analyze_type_info(typ)
+    assert result.kind == "Vector"
+    assert result.elem_type == float
+    assert result.vector_info == VectorInfo(dim=3)
+
+
+def test_ndarray_any_dtype() -> None:
+    typ = NDArray[Any]
+    with pytest.raises(
+        TypeError, match="NDArray for Vector must use a concrete numpy dtype"
+    ):
+        analyze_type_info(typ)
+
+
+def test_list_of_primitives() -> None:
+    typ = List[str]
+    result = analyze_type_info(typ)
+    assert result == AnalyzedTypeInfo(
+        kind="Vector",
+        vector_info=VectorInfo(dim=None),
+        elem_type=str,
+        key_type=None,
+        struct_type=None,
+        np_number_type=None,
+        attrs=None,
+        nullable=False,
+    )
+
+
+def test_list_of_structs() -> None:
+    typ = List[SimpleDataclass]
+    result = analyze_type_info(typ)
+    assert result == AnalyzedTypeInfo(
+        kind="LTable",
+        vector_info=None,
+        elem_type=SimpleDataclass,
+        key_type=None,
+        struct_type=None,
+        np_number_type=None,
+        attrs=None,
+        nullable=False,
+    )
+
+
+def test_sequence_of_int() -> None:
+    typ = Sequence[int]
+    result = analyze_type_info(typ)
+    assert result == AnalyzedTypeInfo(
+        kind="Vector",
+        vector_info=VectorInfo(dim=None),
+        elem_type=int,
+        key_type=None,
+        struct_type=None,
+        np_number_type=None,
+        attrs=None,
+        nullable=False,
+    )
+
+
+def test_list_with_vector_info() -> None:
+    typ = Annotated[List[int], VectorInfo(dim=5)]
+    result = analyze_type_info(typ)
+    assert result == AnalyzedTypeInfo(
+        kind="Vector",
+        vector_info=VectorInfo(dim=5),
+        elem_type=int,
+        key_type=None,
+        struct_type=None,
+        np_number_type=None,
+        attrs=None,
+        nullable=False,
+    )
+
+
+def test_dict_str_int() -> None:
+    typ = Dict[str, int]
+    result = analyze_type_info(typ)
+    assert result == AnalyzedTypeInfo(
+        kind="KTable",
+        vector_info=None,
+        elem_type=(str, int),
+        key_type=None,
+        struct_type=None,
+        np_number_type=None,
+        attrs=None,
+        nullable=False,
+    )
+
+
+def test_mapping_str_dataclass() -> None:
+    typ = Mapping[str, SimpleDataclass]
+    result = analyze_type_info(typ)
+    assert result == AnalyzedTypeInfo(
+        kind="KTable",
+        vector_info=None,
+        elem_type=(str, SimpleDataclass),
+        key_type=None,
+        struct_type=None,
+        np_number_type=None,
+        attrs=None,
+        nullable=False,
+    )
+
+
+def test_dataclass() -> None:
+    typ = SimpleDataclass
+    result = analyze_type_info(typ)
+    assert result == AnalyzedTypeInfo(
+        kind="Struct",
+        vector_info=None,
+        elem_type=None,
+        key_type=None,
+        struct_type=SimpleDataclass,
+        np_number_type=None,
+        attrs=None,
+        nullable=False,
+    )
+
+
+def test_named_tuple() -> None:
+    typ = SimpleNamedTuple
+    result = analyze_type_info(typ)
+    assert result == AnalyzedTypeInfo(
+        kind="Struct",
+        vector_info=None,
+        elem_type=None,
+        key_type=None,
+        struct_type=SimpleNamedTuple,
+        np_number_type=None,
+        attrs=None,
+        nullable=False,
+    )
+
+
+def test_tuple_key_value() -> None:
+    typ = (str, int)
+    result = analyze_type_info(typ)
+    assert result == AnalyzedTypeInfo(
+        kind="Int64",
+        vector_info=None,
+        elem_type=None,
+        key_type=str,
+        struct_type=None,
+        np_number_type=None,
+        attrs=None,
+        nullable=False,
+    )
+
+
+def test_str() -> None:
+    typ = str
+    result = analyze_type_info(typ)
+    assert result == AnalyzedTypeInfo(
+        kind="Str",
+        vector_info=None,
+        elem_type=None,
+        key_type=None,
+        struct_type=None,
+        np_number_type=None,
+        attrs=None,
+        nullable=False,
+    )
+
+
+def test_bool() -> None:
+    typ = bool
+    result = analyze_type_info(typ)
+    assert result == AnalyzedTypeInfo(
+        kind="Bool",
+        vector_info=None,
+        elem_type=None,
+        key_type=None,
+        struct_type=None,
+        np_number_type=None,
+        attrs=None,
+        nullable=False,
+    )
+
+
+def test_bytes() -> None:
+    typ = bytes
+    result = analyze_type_info(typ)
+    assert result == AnalyzedTypeInfo(
+        kind="Bytes",
+        vector_info=None,
+        elem_type=None,
+        key_type=None,
+        struct_type=None,
+        np_number_type=None,
+        attrs=None,
+        nullable=False,
+    )
+
+
+def test_uuid() -> None:
+    typ = uuid.UUID
+    result = analyze_type_info(typ)
+    assert result == AnalyzedTypeInfo(
+        kind="Uuid",
+        vector_info=None,
+        elem_type=None,
+        key_type=None,
+        struct_type=None,
+        np_number_type=None,
+        attrs=None,
+        nullable=False,
+    )
+
+
+def test_date() -> None:
+    typ = datetime.date
+    result = analyze_type_info(typ)
+    assert result == AnalyzedTypeInfo(
+        kind="Date",
+        vector_info=None,
+        elem_type=None,
+        key_type=None,
+        struct_type=None,
+        np_number_type=None,
+        attrs=None,
+        nullable=False,
+    )
+
+
+def test_time() -> None:
+    typ = datetime.time
+    result = analyze_type_info(typ)
+    assert result == AnalyzedTypeInfo(
+        kind="Time",
+        vector_info=None,
+        elem_type=None,
+        key_type=None,
+        struct_type=None,
+        np_number_type=None,
+        attrs=None,
+        nullable=False,
+    )
+
+
+def test_timedelta() -> None:
+    typ = datetime.timedelta
+    result = analyze_type_info(typ)
+    assert result == AnalyzedTypeInfo(
+        kind="TimeDelta",
+        vector_info=None,
+        elem_type=None,
+        key_type=None,
+        struct_type=None,
+        np_number_type=None,
+        attrs=None,
+        nullable=False,
+    )
+
+
+def test_float() -> None:
+    typ = float
+    result = analyze_type_info(typ)
+    assert result == AnalyzedTypeInfo(
+        kind="Float64",
+        vector_info=None,
+        elem_type=None,
+        key_type=None,
+        struct_type=None,
+        np_number_type=None,
+        attrs=None,
+        nullable=False,
+    )
+
+
+def test_int() -> None:
+    typ = int
+    result = analyze_type_info(typ)
+    assert result == AnalyzedTypeInfo(
+        kind="Int64",
+        vector_info=None,
+        elem_type=None,
+        key_type=None,
+        struct_type=None,
+        np_number_type=None,
+        attrs=None,
+        nullable=False,
+    )
+
+
+def test_type_with_attributes() -> None:
+    typ = Annotated[str, TypeAttr("key", "value")]
+    result = analyze_type_info(typ)
+    assert result == AnalyzedTypeInfo(
+        kind="Str",
+        vector_info=None,
+        elem_type=None,
+        key_type=None,
+        struct_type=None,
+        np_number_type=None,
+        attrs={"key": "value"},
+        nullable=False,
+    )
+
+
+def test_encode_enriched_type_none() -> None:
+    typ = None
+    result = encode_enriched_type(typ)
+    assert result is None
+
+
+def test_encode_enriched_type_struct() -> None:
+    typ = SimpleDataclass
+    result = encode_enriched_type(typ)
+    assert result["type"]["kind"] == "Struct"
+    assert len(result["type"]["fields"]) == 2
+    assert result["type"]["fields"][0]["name"] == "name"
+    assert result["type"]["fields"][0]["type"]["kind"] == "Str"
+    assert result["type"]["fields"][1]["name"] == "value"
+    assert result["type"]["fields"][1]["type"]["kind"] == "Int64"
+
+
+def test_encode_enriched_type_vector() -> None:
+    typ = NDArray[np.float32]
+    result = encode_enriched_type(typ)
+    assert result["type"]["kind"] == "Vector"
+    assert result["type"]["element_type"]["kind"] == "Float32"
+    assert result["type"]["dimension"] is None
+
+
+def test_encode_enriched_type_ltable() -> None:
+    typ = List[SimpleDataclass]
+    result = encode_enriched_type(typ)
+    assert result["type"]["kind"] == "LTable"
+    assert result["type"]["row"]["kind"] == "Struct"
+    assert len(result["type"]["row"]["fields"]) == 2
+
+
+def test_encode_enriched_type_with_attrs() -> None:
+    typ = Annotated[str, TypeAttr("key", "value")]
+    result = encode_enriched_type(typ)
+    assert result["type"]["kind"] == "Str"
+    assert result["attrs"] == {"key": "value"}
+
+
+def test_encode_enriched_type_nullable() -> None:
+    typ = str | None
+    result = encode_enriched_type(typ)
+    assert result["type"]["kind"] == "Str"
+    assert result["nullable"] is True
+
+
+def test_invalid_struct_kind() -> None:
+    typ = Annotated[SimpleDataclass, TypeKind("Vector")]
+    with pytest.raises(ValueError, match="Unexpected type kind for struct: Vector"):
+        analyze_type_info(typ)
+
+
+def test_invalid_list_kind() -> None:
+    typ = Annotated[List[int], TypeKind("Struct")]
+    with pytest.raises(ValueError, match="Unexpected type kind for list: Struct"):
+        analyze_type_info(typ)
+
+
+def test_unsupported_type() -> None:
+    typ = set
+    with pytest.raises(ValueError, match="type unsupported yet: <class 'set'>"):
+        analyze_type_info(typ)

cocoindex/typing.py

@@ -9,14 +9,16 @@ from typing import (
     Annotated,
     NamedTuple,
     Any,
+    KeysView,
     TypeVar,
     TYPE_CHECKING,
     overload,
-    Sequence,
     Generic,
     Literal,
     Protocol,
 )
+import numpy as np
+from numpy.typing import NDArray
 
 
 class VectorInfo(NamedTuple):
@@ -47,10 +49,10 @@ OffsetDateTime = Annotated[datetime.datetime, TypeKind("OffsetDateTime")]
 
 if TYPE_CHECKING:
     T_co = TypeVar("T_co", covariant=True)
-    Dim_co = TypeVar("Dim_co", bound=int, covariant=True)
+    Dim_co = TypeVar("Dim_co", bound=int | None, covariant=True, default=None)
 
     class Vector(Protocol, Generic[T_co, Dim_co]):
-        """Vector[T, Dim] is a special typing alias for a list[T] with optional dimension info"""
+        """Vector[T, Dim] is a special typing alias for an NDArray[T] with optional dimension info"""
 
         def __getitem__(self, index: int) -> T_co: ...
         def __len__(self) -> int: ...
@@ -58,25 +60,34 @@ if TYPE_CHECKING:
 else:
 
     class Vector:  # type: ignore[unreachable]
-        """A special typing alias for a list[T] with optional dimension info"""
+        """A special typing alias for an NDArray[T] with optional dimension info"""
 
         def __class_getitem__(self, params):
             if not isinstance(params, tuple):
-                # Only element type provided
-                elem_type = params
-                return Annotated[list[elem_type], VectorInfo(dim=None)]
+                # No dimension provided, e.g., Vector[np.float32]
+                dtype = params
+                # Use NDArray for supported numeric dtypes, else list
+                if DtypeRegistry.get_by_dtype(dtype) is not None:
+                    return Annotated[NDArray[dtype], VectorInfo(dim=None)]
+                return Annotated[list[dtype], VectorInfo(dim=None)]
             else:
-                # Element type and dimension provided
-                elem_type, dim = params
-                if typing.get_origin(dim) is Literal:
-                    dim = typing.get_args(dim)[0]  # Extract the literal value
-                return Annotated[list[elem_type], VectorInfo(dim=dim)]
+                # Element type and dimension provided, e.g., Vector[np.float32, Literal[3]]
+                dtype, dim_literal = params
+                # Extract the literal value
+                dim_val = (
+                    typing.get_args(dim_literal)[0]
+                    if typing.get_origin(dim_literal) is Literal
+                    else None
+                )
+                if DtypeRegistry.get_by_dtype(dtype) is not None:
+                    return Annotated[NDArray[dtype], VectorInfo(dim=dim_val)]
+                return Annotated[list[dtype], VectorInfo(dim=dim_val)]
 
 
 TABLE_TYPES: tuple[str, str] = ("KTable", "LTable")
 KEY_FIELD_NAME: str = "_key"
 
-ElementType = type | tuple[type, type]
+ElementType = type | tuple[type, type] | Annotated[Any, TypeKind]
 
 
 def is_namedtuple_type(t: type) -> bool:
@@ -89,6 +100,43 @@ def _is_struct_type(t: ElementType | None) -> bool:
     )
 
 
+class DtypeInfo:
+    """Metadata for a NumPy dtype."""
+
+    def __init__(self, numpy_dtype: type, kind: str, python_type: type) -> None:
+        self.numpy_dtype = numpy_dtype
+        self.kind = kind
+        self.python_type = python_type
+        self.annotated_type = Annotated[python_type, TypeKind(kind)]
+
+
+class DtypeRegistry:
+    """
+    Registry for NumPy dtypes used in CocoIndex.
+    Provides mappings from NumPy dtypes to CocoIndex's type representation.
+    """
+
+    _mappings: dict[type, DtypeInfo] = {
+        np.float32: DtypeInfo(np.float32, "Float32", float),
+        np.float64: DtypeInfo(np.float64, "Float64", float),
+        np.int64: DtypeInfo(np.int64, "Int64", int),
+    }
+
+    @classmethod
+    def get_by_dtype(cls, dtype: Any) -> DtypeInfo | None:
+        """Get DtypeInfo by NumPy dtype."""
+        if dtype is Any:
+            raise TypeError(
+                "NDArray for Vector must use a concrete numpy dtype, got `Any`."
+            )
+        return cls._mappings.get(dtype)
+
+    @staticmethod
+    def supported_dtypes() -> KeysView[type]:
+        """Get a list of supported NumPy dtypes."""
+        return DtypeRegistry._mappings.keys()
+
+
 @dataclasses.dataclass
 class AnalyzedTypeInfo:
     """
@@ -101,6 +149,9 @@ class AnalyzedTypeInfo:
 
     key_type: type | None  # For element of KTable
     struct_type: type | None  # For Struct, a dataclass or namedtuple
+    np_number_type: (
+        type | None
+    )  # NumPy dtype for the element type, if represented by numpy.ndarray or a NumPy scalar
 
     attrs: dict[str, Any] | None
     nullable: bool = False
@@ -155,6 +206,7 @@ def analyze_type_info(t: Any) -> AnalyzedTypeInfo:
     struct_type: type | None = None
     elem_type: ElementType | None = None
     key_type: type | None = None
+    np_number_type: type | None = None
     if _is_struct_type(t):
         struct_type = t
 
@@ -179,12 +231,35 @@ def analyze_type_info(t: Any) -> AnalyzedTypeInfo:
                     vector_info = VectorInfo(dim=None)
         elif not (kind == "Vector" or kind in TABLE_TYPES):
             raise ValueError(f"Unexpected type kind for list: {kind}")
+    elif base_type is np.ndarray:
+        kind = "Vector"
+        args = typing.get_args(t)
+        _, dtype_spec = args
+
+        dtype_args = typing.get_args(dtype_spec)
+        if not dtype_args:
+            raise ValueError("Invalid dtype specification for NDArray")
+
+        np_number_type = dtype_args[0]
+        dtype_info = DtypeRegistry.get_by_dtype(np_number_type)
+        if dtype_info is None:
+            raise ValueError(
+                f"Unsupported numpy dtype for NDArray: {np_number_type}. "
+                f"Supported dtypes: {DtypeRegistry.supported_dtypes()}"
+            )
+        elem_type = dtype_info.annotated_type
+        vector_info = VectorInfo(dim=None) if vector_info is None else vector_info
+
     elif base_type is collections.abc.Mapping or base_type is dict:
         args = typing.get_args(t)
         elem_type = (args[0], args[1])
         kind = "KTable"
     elif kind is None:
-        if t is bytes:
+        dtype_info = DtypeRegistry.get_by_dtype(t)
+        if dtype_info is not None:
+            kind = dtype_info.kind
+            np_number_type = dtype_info.numpy_dtype
+        elif t is bytes:
             kind = "Bytes"
         elif t is str:
             kind = "Str"
@@ -213,6 +288,7 @@ def analyze_type_info(t: Any) -> AnalyzedTypeInfo:
         elem_type=elem_type,
         key_type=key_type,
         struct_type=struct_type,
+        np_number_type=np_number_type,
         attrs=attrs,
         nullable=nullable,
     )
@@ -265,9 +341,8 @@ def _encode_type(type_info: AnalyzedTypeInfo) -> dict[str, Any]:
             raise ValueError("Vector type must have a vector info")
         if type_info.elem_type is None:
             raise ValueError("Vector type must have an element type")
-        encoded_type["element_type"] = _encode_type(
-            analyze_type_info(type_info.elem_type)
-        )
+        elem_type_info = analyze_type_info(type_info.elem_type)
+        encoded_type["element_type"] = _encode_type(elem_type_info)
         encoded_type["dimension"] = type_info.vector_info.dim
 
     elif type_info.kind in TABLE_TYPES:

{cocoindex-0.1.49.dist-info → cocoindex-0.1.51.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: cocoindex
-Version: 0.1.49
+Version: 0.1.51
 Requires-Dist: sentence-transformers>=3.3.1
 Requires-Dist: click>=8.1.8
 Requires-Dist: rich>=14.0.0