cocoindex 0.1.52__cp311-cp311-win_amd64.whl → 0.1.54__cp311-cp311-win_amd64.whl

cocoindex/tests/test_typing.py

@@ -1,30 +1,21 @@
 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
+from collections.abc import Mapping, Sequence
+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 (
-    analyze_type_info,
+    AnalyzedTypeInfo,
+    TypeAttr,
+    TypeKind,
     Vector,
     VectorInfo,
-    TypeKind,
-    TypeAttr,
-    Float32,
-    Float64,
+    analyze_type_info,
     encode_enriched_type,
-    AnalyzedTypeInfo,
 )
 
 
@@ -42,61 +33,57 @@ class SimpleNamedTuple(NamedTuple):
 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,
-    )
+    assert result.kind == "Vector"
+    assert result.vector_info == VectorInfo(dim=None)
+    assert result.elem_type == np.float32
+    assert result.key_type is None
+    assert result.struct_type is None
+    assert result.nullable is False
+    assert result.np_number_type is not None
+    assert get_origin(result.np_number_type) == np.ndarray
+    assert get_args(result.np_number_type)[1] == np.dtype[np.float32]
 
 
 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,
-    )
+    assert result.kind == "Vector"
+    assert result.vector_info == VectorInfo(dim=None)
+    assert result.elem_type == np.float32
+    assert result.key_type is None
+    assert result.struct_type is None
+    assert result.nullable is False
+    assert result.np_number_type is not None
+    assert get_origin(result.np_number_type) == np.ndarray
+    assert get_args(result.np_number_type)[1] == np.dtype[np.float32]
 
 
 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,
-    )
+    assert result.kind == "Vector"
+    assert result.vector_info == VectorInfo(dim=128)
+    assert result.elem_type == np.float64
+    assert result.key_type is None
+    assert result.struct_type is None
+    assert result.nullable is False
+    assert result.np_number_type is not None
+    assert get_origin(result.np_number_type) == np.ndarray
+    assert get_args(result.np_number_type)[1] == np.dtype[np.float64]
 
 
 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,
-    )
+    assert result.kind == "Vector"
+    assert result.vector_info == VectorInfo(dim=384)
+    assert result.elem_type == np.float32
+    assert result.key_type is None
+    assert result.struct_type is None
+    assert result.nullable is False
+    assert result.np_number_type is not None
+    assert get_origin(result.np_number_type) == np.ndarray
+    assert get_args(result.np_number_type)[1] == np.dtype[np.float32]
 
 
 def test_ndarray_int64_no_dim() -> None:
@@ -104,30 +91,49 @@ def test_ndarray_int64_no_dim() -> None:
     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
+    assert result.elem_type == np.int64
+    assert result.nullable is False
+    assert result.np_number_type is not None
+    assert get_origin(result.np_number_type) == np.ndarray
+    assert get_args(result.np_number_type)[1] == np.dtype[np.int64]
 
 
 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,
-    )
+    assert result.kind == "Vector"
+    assert result.vector_info == VectorInfo(dim=None)
+    assert result.elem_type == np.float32
+    assert result.key_type is None
+    assert result.struct_type is None
+    assert result.nullable is True
+    assert result.np_number_type is not None
+    assert get_origin(result.np_number_type) == np.ndarray
+    assert get_args(result.np_number_type)[1] == np.dtype[np.float32]
+
+
+def test_scalar_numpy_types() -> None:
+    for np_type, expected_kind in [
+        (np.int64, "Int64"),
+        (np.float32, "Float32"),
+        (np.float64, "Float64"),
+    ]:
+        type_info = analyze_type_info(np_type)
+        assert type_info.kind == expected_kind, (
+            f"Expected {expected_kind} for {np_type}, got {type_info.kind}"
+        )
+        assert type_info.np_number_type == np_type, (
+            f"Expected {np_type}, got {type_info.np_number_type}"
+        )
+        assert type_info.elem_type is None
+        assert type_info.vector_info is None
 
 
 def test_vector_str() -> None:
     typ = Vector[str]
     result = analyze_type_info(typ)
     assert result.kind == "Vector"
-    assert result.elem_type == str
+    assert result.elem_type is str
     assert result.vector_info == VectorInfo(dim=None)
 
 
@@ -143,7 +149,7 @@ 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.elem_type is float
     assert result.vector_info == VectorInfo(dim=3)
 
 
@@ -156,10 +162,11 @@ def test_ndarray_any_dtype() -> None:
 
 
 def test_list_of_primitives() -> None:
-    typ = List[str]
+    typ = list[str]
     result = analyze_type_info(typ)
     assert result == AnalyzedTypeInfo(
         kind="Vector",
+        core_type=list[str],
         vector_info=VectorInfo(dim=None),
         elem_type=str,
         key_type=None,
@@ -171,10 +178,11 @@ def test_list_of_primitives() -> None:
 
 
 def test_list_of_structs() -> None:
-    typ = List[SimpleDataclass]
+    typ = list[SimpleDataclass]
     result = analyze_type_info(typ)
     assert result == AnalyzedTypeInfo(
         kind="LTable",
+        core_type=list[SimpleDataclass],
         vector_info=None,
         elem_type=SimpleDataclass,
         key_type=None,
@@ -190,6 +198,7 @@ def test_sequence_of_int() -> None:
     result = analyze_type_info(typ)
     assert result == AnalyzedTypeInfo(
         kind="Vector",
+        core_type=Sequence[int],
         vector_info=VectorInfo(dim=None),
         elem_type=int,
         key_type=None,
@@ -201,10 +210,11 @@ def test_sequence_of_int() -> None:
 
 
 def test_list_with_vector_info() -> None:
-    typ = Annotated[List[int], VectorInfo(dim=5)]
+    typ = Annotated[list[int], VectorInfo(dim=5)]
     result = analyze_type_info(typ)
     assert result == AnalyzedTypeInfo(
         kind="Vector",
+        core_type=list[int],
         vector_info=VectorInfo(dim=5),
         elem_type=int,
         key_type=None,
@@ -216,10 +226,11 @@ def test_list_with_vector_info() -> None:
 
 
 def test_dict_str_int() -> None:
-    typ = Dict[str, int]
+    typ = dict[str, int]
     result = analyze_type_info(typ)
     assert result == AnalyzedTypeInfo(
         kind="KTable",
+        core_type=dict[str, int],
         vector_info=None,
         elem_type=(str, int),
         key_type=None,
@@ -235,6 +246,7 @@ def test_mapping_str_dataclass() -> None:
     result = analyze_type_info(typ)
     assert result == AnalyzedTypeInfo(
         kind="KTable",
+        core_type=Mapping[str, SimpleDataclass],
         vector_info=None,
         elem_type=(str, SimpleDataclass),
         key_type=None,
@@ -250,6 +262,7 @@ def test_dataclass() -> None:
     result = analyze_type_info(typ)
     assert result == AnalyzedTypeInfo(
         kind="Struct",
+        core_type=SimpleDataclass,
         vector_info=None,
         elem_type=None,
         key_type=None,
@@ -265,6 +278,7 @@ def test_named_tuple() -> None:
     result = analyze_type_info(typ)
     assert result == AnalyzedTypeInfo(
         kind="Struct",
+        core_type=SimpleNamedTuple,
         vector_info=None,
         elem_type=None,
         key_type=None,
@@ -280,6 +294,7 @@ def test_tuple_key_value() -> None:
     result = analyze_type_info(typ)
     assert result == AnalyzedTypeInfo(
         kind="Int64",
+        core_type=int,
         vector_info=None,
         elem_type=None,
         key_type=str,
@@ -295,6 +310,7 @@ def test_str() -> None:
     result = analyze_type_info(typ)
     assert result == AnalyzedTypeInfo(
         kind="Str",
+        core_type=str,
         vector_info=None,
         elem_type=None,
         key_type=None,
@@ -310,6 +326,7 @@ def test_bool() -> None:
     result = analyze_type_info(typ)
     assert result == AnalyzedTypeInfo(
         kind="Bool",
+        core_type=bool,
         vector_info=None,
         elem_type=None,
         key_type=None,
@@ -325,6 +342,7 @@ def test_bytes() -> None:
     result = analyze_type_info(typ)
     assert result == AnalyzedTypeInfo(
         kind="Bytes",
+        core_type=bytes,
         vector_info=None,
         elem_type=None,
         key_type=None,
@@ -340,6 +358,7 @@ def test_uuid() -> None:
     result = analyze_type_info(typ)
     assert result == AnalyzedTypeInfo(
         kind="Uuid",
+        core_type=uuid.UUID,
         vector_info=None,
         elem_type=None,
         key_type=None,
@@ -355,6 +374,7 @@ def test_date() -> None:
     result = analyze_type_info(typ)
     assert result == AnalyzedTypeInfo(
         kind="Date",
+        core_type=datetime.date,
         vector_info=None,
         elem_type=None,
         key_type=None,
@@ -370,6 +390,7 @@ def test_time() -> None:
     result = analyze_type_info(typ)
     assert result == AnalyzedTypeInfo(
         kind="Time",
+        core_type=datetime.time,
         vector_info=None,
         elem_type=None,
         key_type=None,
@@ -385,6 +406,7 @@ def test_timedelta() -> None:
     result = analyze_type_info(typ)
     assert result == AnalyzedTypeInfo(
         kind="TimeDelta",
+        core_type=datetime.timedelta,
         vector_info=None,
         elem_type=None,
         key_type=None,
@@ -400,6 +422,7 @@ def test_float() -> None:
     result = analyze_type_info(typ)
     assert result == AnalyzedTypeInfo(
         kind="Float64",
+        core_type=float,
         vector_info=None,
         elem_type=None,
         key_type=None,
@@ -415,6 +438,7 @@ def test_int() -> None:
     result = analyze_type_info(typ)
     assert result == AnalyzedTypeInfo(
         kind="Int64",
+        core_type=int,
         vector_info=None,
         elem_type=None,
         key_type=None,
@@ -430,6 +454,7 @@ def test_type_with_attributes() -> None:
     result = analyze_type_info(typ)
     assert result == AnalyzedTypeInfo(
         kind="Str",
+        core_type=str,
         vector_info=None,
         elem_type=None,
         key_type=None,
@@ -466,7 +491,7 @@ def test_encode_enriched_type_vector() -> None:
 
 
 def test_encode_enriched_type_ltable() -> None:
-    typ = List[SimpleDataclass]
+    typ = list[SimpleDataclass]
     result = encode_enriched_type(typ)
     assert result["type"]["kind"] == "LTable"
     assert result["type"]["row"]["kind"] == "Struct"
@@ -487,6 +512,19 @@ def test_encode_enriched_type_nullable() -> None:
     assert result["nullable"] is True
 
 
+def test_encode_scalar_numpy_types_schema() -> None:
+    for np_type, expected_kind in [
+        (np.int64, "Int64"),
+        (np.float32, "Float32"),
+        (np.float64, "Float64"),
+    ]:
+        schema = encode_enriched_type(np_type)
+        assert schema["type"]["kind"] == expected_kind, (
+            f"Expected {expected_kind} for {np_type}, got {schema['type']['kind']}"
+        )
+        assert not schema.get("nullable", False)
+
+
 def test_invalid_struct_kind() -> None:
     typ = Annotated[SimpleDataclass, TypeKind("Vector")]
     with pytest.raises(ValueError, match="Unexpected type kind for struct: Vector"):
@@ -494,7 +532,7 @@ def test_invalid_struct_kind() -> None:
 
 
 def test_invalid_list_kind() -> None:
-    typ = Annotated[List[int], TypeKind("Struct")]
+    typ = Annotated[list[int], TypeKind("Struct")]
     with pytest.raises(ValueError, match="Unexpected type kind for list: Struct"):
         analyze_type_info(typ)
 

cocoindex/typing.py

@@ -1,22 +1,22 @@
-import typing
 import collections
 import dataclasses
 import datetime
-import types
 import inspect
+import types
+import typing
 import uuid
 from typing import (
+    TYPE_CHECKING,
     Annotated,
-    NamedTuple,
     Any,
-    KeysView,
-    TypeVar,
-    TYPE_CHECKING,
-    overload,
     Generic,
     Literal,
+    NamedTuple,
     Protocol,
+    TypeVar,
+    overload,
 )
+
 import numpy as np
 from numpy.typing import NDArray
 
@@ -67,7 +67,7 @@ else:
                 # 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:
+                if dtype in DtypeRegistry._DTYPE_TO_KIND:
                     return Annotated[NDArray[dtype], VectorInfo(dim=None)]
                 return Annotated[list[dtype], VectorInfo(dim=None)]
             else:
@@ -79,7 +79,7 @@ else:
                     if typing.get_origin(dim_literal) is Literal
                     else None
                 )
-                if DtypeRegistry.get_by_dtype(dtype) is not None:
+                if dtype in DtypeRegistry._DTYPE_TO_KIND:
                     return Annotated[NDArray[dtype], VectorInfo(dim=dim_val)]
                 return Annotated[list[dtype], VectorInfo(dim=dim_val)]
 
@@ -90,6 +90,19 @@ KEY_FIELD_NAME: str = "_key"
 ElementType = type | tuple[type, type] | Annotated[Any, TypeKind]
 
 
+def extract_ndarray_scalar_dtype(ndarray_type: Any) -> Any:
+    args = typing.get_args(ndarray_type)
+    _, dtype_spec = args
+    dtype_args = typing.get_args(dtype_spec)
+    if not dtype_args:
+        raise ValueError(f"Invalid dtype specification: {dtype_spec}")
+    return dtype_args[0]
+
+
+def is_numpy_number_type(t: type) -> bool:
+    return isinstance(t, type) and issubclass(t, np.number)
+
+
 def is_namedtuple_type(t: type) -> bool:
     return isinstance(t, type) and issubclass(t, tuple) and hasattr(t, "_fields")
 
@@ -100,41 +113,34 @@ 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.
+    Maps NumPy dtypes to their CocoIndex type kind.
     """
 
-    _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),
+    _DTYPE_TO_KIND: dict[ElementType, str] = {
+        np.float32: "Float32",
+        np.float64: "Float64",
+        np.int64: "Int64",
     }
 
     @classmethod
-    def get_by_dtype(cls, dtype: Any) -> DtypeInfo | None:
-        """Get DtypeInfo by NumPy dtype."""
+    def validate_dtype_and_get_kind(cls, dtype: ElementType) -> str:
+        """
+        Validate that the given dtype is supported, and get its CocoIndex kind by 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()
+        kind = cls._DTYPE_TO_KIND.get(dtype)
+        if kind is None:
+            raise ValueError(
+                f"Unsupported NumPy dtype in NDArray: {dtype}. "
+                f"Supported dtypes: {cls._DTYPE_TO_KIND.keys()}"
+            )
+        return kind
 
 
 @dataclasses.dataclass
@@ -144,6 +150,7 @@ class AnalyzedTypeInfo:
     """
 
     kind: str
+    core_type: Any
     vector_info: VectorInfo | None  # For Vector
     elem_type: ElementType | None  # For Vector and Table
 
@@ -155,6 +162,7 @@ class AnalyzedTypeInfo:
 
     attrs: dict[str, Any] | None
     nullable: bool = False
+    union_variant_types: typing.List[ElementType] | None = None  # For Union
 
 
 def analyze_type_info(t: Any) -> AnalyzedTypeInfo:
@@ -175,18 +183,6 @@ def analyze_type_info(t: Any) -> AnalyzedTypeInfo:
         if base_type is Annotated:
             annotations = t.__metadata__
             t = t.__origin__
-        elif base_type is types.UnionType:
-            possible_types = typing.get_args(t)
-            non_none_types = [
-                arg for arg in possible_types if arg not in (None, types.NoneType)
-            ]
-            if len(non_none_types) != 1:
-                raise ValueError(
-                    f"Expect exactly one non-None choice for Union type, but got {len(non_none_types)}: {t}"
-                )
-            t = non_none_types[0]
-            if len(possible_types) > 1:
-                nullable = True
         else:
             break
 
@@ -205,6 +201,7 @@ def analyze_type_info(t: Any) -> AnalyzedTypeInfo:
 
     struct_type: type | None = None
     elem_type: ElementType | None = None
+    union_variant_types: typing.List[ElementType] | None = None
     key_type: type | None = None
     np_number_type: type | None = None
     if _is_struct_type(t):
@@ -214,6 +211,9 @@ def analyze_type_info(t: Any) -> AnalyzedTypeInfo:
             kind = "Struct"
         elif kind != "Struct":
             raise ValueError(f"Unexpected type kind for struct: {kind}")
+    elif is_numpy_number_type(t):
+        np_number_type = t
+        kind = DtypeRegistry.validate_dtype_and_get_kind(t)
     elif base_type is collections.abc.Sequence or base_type is list:
         args = typing.get_args(t)
         elem_type = args[0]
@@ -233,33 +233,35 @@ def analyze_type_info(t: Any) -> AnalyzedTypeInfo:
             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
+        np_number_type = t
+        elem_type = extract_ndarray_scalar_dtype(np_number_type)
+        _ = DtypeRegistry.validate_dtype_and_get_kind(elem_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 base_type is types.UnionType:
+        possible_types = typing.get_args(t)
+        non_none_types = [
+            arg for arg in possible_types if arg not in (None, types.NoneType)
+        ]
+
+        if len(non_none_types) == 0:
+            return analyze_type_info(None)
+
+        nullable = len(non_none_types) < len(possible_types)
+
+        if len(non_none_types) == 1:
+            result = analyze_type_info(non_none_types[0])
+            result.nullable = nullable
+            return result
+
+        kind = "Union"
+        union_variant_types = non_none_types
     elif kind is None:
-        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:
+        if t is bytes:
             kind = "Bytes"
         elif t is str:
             kind = "Str"
@@ -284,8 +286,10 @@ def analyze_type_info(t: Any) -> AnalyzedTypeInfo:
 
     return AnalyzedTypeInfo(
         kind=kind,
+        core_type=t,
         vector_info=vector_info,
         elem_type=elem_type,
+        union_variant_types=union_variant_types,
         key_type=key_type,
         struct_type=struct_type,
         np_number_type=np_number_type,
@@ -345,6 +349,14 @@ def _encode_type(type_info: AnalyzedTypeInfo) -> dict[str, Any]:
         encoded_type["element_type"] = _encode_type(elem_type_info)
         encoded_type["dimension"] = type_info.vector_info.dim
 
+    elif type_info.kind == "Union":
+        if type_info.union_variant_types is None:
+            raise ValueError("Union type must have a variant type list")
+        encoded_type["types"] = [
+            _encode_type(analyze_type_info(typ))
+            for typ in type_info.union_variant_types
+        ]
+
     elif type_info.kind in TABLE_TYPES:
         if type_info.elem_type is None:
             raise ValueError(f"{type_info.kind} type must have an element type")