refactor: simplify DtypeRegistry structure and associated methods

Author: lemorage

python/cocoindex/convert.py

@@ -6,18 +6,19 @@
 import datetime
 import inspect
 import uuid
+from enum import Enum
+from typing import Any, Callable, Mapping, get_origin
+
 import numpy as np
 
-from enum import Enum
-from typing import Any, Callable, get_origin, Mapping
 from .typing import (
+    KEY_FIELD_NAME,
+    TABLE_TYPES,
+    DtypeRegistry,
     analyze_type_info,
     encode_enriched_type,
     extract_ndarray_scalar_dtype,
     is_namedtuple_type,
-    TABLE_TYPES,
-    KEY_FIELD_NAME,
-    DtypeRegistry,
 )
 
 
@@ -178,12 +179,7 @@ def decode_vector(value: Any) -> Any | None:
                 scalar_dtype = extract_ndarray_scalar_dtype(
                     dst_type_info.np_number_type
                 )
-                dtype_info = DtypeRegistry.get_by_dtype(scalar_dtype)
-                if dtype_info is None:
-                    raise ValueError(
-                        f"Unsupported dtype in NDArray: {scalar_dtype}. "
-                        f"Supported dtypes: {DtypeRegistry.supported_dtypes()}"
-                    )
+                _ = DtypeRegistry.validate_and_get_dtype_info(scalar_dtype)
                 return np.array(value, dtype=scalar_dtype)
 
         return decode_vector

python/cocoindex/tests/test_typing.py

@@ -10,8 +10,6 @@
 
 from cocoindex.typing import (
     AnalyzedTypeInfo,
-    Float32,
-    Float64,
     TypeAttr,
     TypeKind,
     Vector,
@@ -37,7 +35,7 @@ def test_ndarray_float32_no_dim() -> None:
     result = analyze_type_info(typ)
     assert result.kind == "Vector"
     assert result.vector_info == VectorInfo(dim=None)
-    assert result.elem_type == Float32
+    assert result.elem_type == np.float32
     assert result.key_type is None
     assert result.struct_type is None
     assert result.nullable is False
@@ -51,7 +49,7 @@ def test_vector_float32_no_dim() -> None:
     result = analyze_type_info(typ)
     assert result.kind == "Vector"
     assert result.vector_info == VectorInfo(dim=None)
-    assert result.elem_type == Float32
+    assert result.elem_type == np.float32
     assert result.key_type is None
     assert result.struct_type is None
     assert result.nullable is False
@@ -65,7 +63,7 @@ def test_ndarray_float64_with_dim() -> None:
     result = analyze_type_info(typ)
     assert result.kind == "Vector"
     assert result.vector_info == VectorInfo(dim=128)
-    assert result.elem_type == Float64
+    assert result.elem_type == np.float64
     assert result.key_type is None
     assert result.struct_type is None
     assert result.nullable is False
@@ -79,7 +77,7 @@ def test_vector_float32_with_dim() -> None:
     result = analyze_type_info(typ)
     assert result.kind == "Vector"
     assert result.vector_info == VectorInfo(dim=384)
-    assert result.elem_type == Float32
+    assert result.elem_type == np.float32
     assert result.key_type is None
     assert result.struct_type is None
     assert result.nullable is False
@@ -93,7 +91,7 @@ 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 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
@@ -105,7 +103,7 @@ def test_nullable_ndarray() -> None:
     result = analyze_type_info(typ)
     assert result.kind == "Vector"
     assert result.vector_info == VectorInfo(dim=None)
-    assert result.elem_type == Float32
+    assert result.elem_type == np.float32
     assert result.key_type is None
     assert result.struct_type is None
     assert result.nullable is True
@@ -121,12 +119,12 @@ def test_scalar_numpy_types() -> None:
         (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.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
 
@@ -502,9 +500,9 @@ def test_encode_scalar_numpy_types_schema() -> None:
         (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 (
+            schema["type"]["kind"] == expected_kind
+        ), f"Expected {expected_kind} for {np_type}, got {schema['type']['kind']}"
         assert not schema.get("nullable", False)
 
 

python/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
 
@@ -113,41 +113,40 @@ 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[type, 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 get_by_dtype(cls, dtype: Any) -> tuple[type, str] | None:
+        """Get the NumPy dtype and 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)
+        kind = cls._DTYPE_TO_KIND.get(dtype)
+        return None if kind is None else (dtype, kind)
 
-    @staticmethod
-    def supported_dtypes() -> KeysView[type]:
-        """Get a list of supported NumPy dtypes."""
-        return DtypeRegistry._mappings.keys()
+    @classmethod
+    def validate_and_get_dtype_info(cls, dtype: Any) -> tuple[type, str]:
+        """
+        Validate that the given dtype is supported.
+        """
+        dtype_info = cls.get_by_dtype(dtype)
+        if dtype_info is None:
+            raise ValueError(
+                f"Unsupported NumPy dtype in NDArray: {dtype}. "
+                f"Supported dtypes: {cls._DTYPE_TO_KIND.keys()}"
+            )
+        return dtype_info
 
 
 @dataclasses.dataclass
@@ -228,11 +227,7 @@ def analyze_type_info(t: Any) -> AnalyzedTypeInfo:
         elif kind != "Struct":
             raise ValueError(f"Unexpected type kind for struct: {kind}")
     elif is_numpy_number_type(t):
-        if (dtype_info := DtypeRegistry.get_by_dtype(t)) is not None:
-            kind = dtype_info.kind
-            np_number_type = dtype_info.numpy_dtype
-        else:
-            raise ValueError(f"Unsupported NumPy dtype: {t}")
+        np_number_type, kind = DtypeRegistry.validate_and_get_dtype_info(t)
     elif base_type is collections.abc.Sequence or base_type is list:
         args = typing.get_args(t)
         elem_type = args[0]
@@ -253,14 +248,8 @@ def analyze_type_info(t: Any) -> AnalyzedTypeInfo:
     elif base_type is np.ndarray:
         kind = "Vector"
         np_number_type = t
-        numpy_dtype = extract_ndarray_scalar_dtype(np_number_type)
-        dtype_info = DtypeRegistry.get_by_dtype(numpy_dtype)
-        if dtype_info is None:
-            raise ValueError(
-                f"Unsupported numpy dtype for NDArray: {numpy_dtype}. "
-                f"Supported dtypes: {DtypeRegistry.supported_dtypes()}"
-            )
-        elem_type = dtype_info.annotated_type
+        elem_type = extract_ndarray_scalar_dtype(np_number_type)
+        _ = DtypeRegistry.validate_and_get_dtype_info(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: