fix: reorder imports and annotate return types

Author: lemorage

python/cocoindex/tests/test_convert.py

@@ -1,24 +1,25 @@
-import uuid
 import datetime
+import uuid
 from dataclasses import dataclass, make_dataclass
-from typing import NamedTuple, Literal, Any, Annotated, Callable
+from typing import Annotated, Any, Callable, Literal, NamedTuple
+
+import numpy as np
 import pytest
+from numpy.typing import NDArray
+
 import cocoindex
+from cocoindex.convert import (
+    dump_engine_object,
+    encode_engine_value,
+    make_engine_value_decoder,
+)
 from cocoindex.typing import (
-    encode_enriched_type,
     Float32,
+    Float64,
     TypeKind,
     Vector,
-    Float32,
-    Float64,
-)
-from cocoindex.convert import (
-    encode_engine_value,
-    make_engine_value_decoder,
-    dump_engine_object,
+    encode_enriched_type,
 )
-import numpy as np
-from numpy.typing import NDArray
 
 
 @dataclass
@@ -130,7 +131,7 @@ def test_encode_engine_value_date_time_types() -> None:
     assert encode_engine_value(dt) == dt
 
 
-def test_encode_scalar_numpy_values():
+def test_encode_scalar_numpy_values() -> None:
     """Test encoding scalar NumPy values to engine-compatible values."""
     test_cases = [
         (np.int64(42), 42),
@@ -228,7 +229,7 @@ def test_roundtrip_basic_types() -> None:
     )
 
 
-def test_decode_scalar_numpy_values():
+def test_decode_scalar_numpy_values() -> None:
     test_cases = [
         ({"kind": "Int64"}, np.int64, 42, np.int64(42)),
         ({"kind": "Float32"}, np.float32, 3.14, np.float32(3.14)),
@@ -241,29 +242,29 @@ def test_decode_scalar_numpy_values():
         assert result == expected
 
 
-def test_non_ndarray_vector_decoding():
+def test_non_ndarray_vector_decoding() -> None:
     # Test list[np.float64]
     src_type = {
         "kind": "Vector",
         "element_type": {"kind": "Float64"},
         "dimension": None,
     }
-    dst_type = list[np.float64]
-    decoder = make_engine_value_decoder(["field"], src_type, dst_type)
-    input_value = [1.0, 2.0, 3.0]
-    result = decoder(input_value)
+    dst_type_float = list[np.float64]
+    decoder = make_engine_value_decoder(["field"], src_type, dst_type_float)
+    input_numbers = [1.0, 2.0, 3.0]
+    result = decoder(input_numbers)
     assert isinstance(result, list)
     assert all(isinstance(x, np.float64) for x in result)
     assert result == [np.float64(1.0), np.float64(2.0), np.float64(3.0)]
 
     # Test list[Uuid]
     src_type = {"kind": "Vector", "element_type": {"kind": "Uuid"}, "dimension": None}
-    dst_type = list[uuid.UUID]
-    decoder = make_engine_value_decoder(["field"], src_type, dst_type)
+    dst_type_uuid = list[uuid.UUID]
+    decoder = make_engine_value_decoder(["field"], src_type, dst_type_uuid)
     uuid1 = uuid.uuid4()
     uuid2 = uuid.uuid4()
-    input_value = [uuid1.bytes, uuid2.bytes]
-    result = decoder(input_value)
+    input_bytes = [uuid1.bytes, uuid2.bytes]
+    result = decoder(input_bytes)
     assert isinstance(result, list)
     assert all(isinstance(x, uuid.UUID) for x in result)
     assert result == [uuid1, uuid2]
@@ -946,6 +947,6 @@ class MixedStruct:
         python_float=3.14,
         string="hello, world",
         annotated_int=np.int64(42),
-        annotated_float=np.float32(3.14),
+        annotated_float=2.0,
     )
     validate_full_roundtrip(instance, MixedStruct)

python/cocoindex/tests/test_typing.py

@@ -1,30 +1,23 @@
 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, Dict, List, Literal, NamedTuple, get_args, get_origin
+
 import numpy as np
+import pytest
 from numpy.typing import NDArray
 
 from cocoindex.typing import (
-    analyze_type_info,
-    Vector,
-    VectorInfo,
-    TypeKind,
-    TypeAttr,
+    AnalyzedTypeInfo,
     Float32,
     Float64,
+    TypeAttr,
+    TypeKind,
+    Vector,
+    VectorInfo,
+    analyze_type_info,
     encode_enriched_type,
-    AnalyzedTypeInfo,
 )
 
 
@@ -42,61 +35,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.ndarray[tuple[int, ...], np.dtype[np.float32]],
-        attrs=None,
-        nullable=False,
-    )
+    assert result.kind == "Vector"
+    assert result.vector_info == VectorInfo(dim=None)
+    assert result.elem_type == 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.ndarray[tuple[int, ...], np.dtype[np.float32]],
-        attrs=None,
-        nullable=False,
-    )
+    assert result.kind == "Vector"
+    assert result.vector_info == VectorInfo(dim=None)
+    assert result.elem_type == 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.ndarray[tuple[int, ...], np.dtype[np.float64]],
-        attrs=None,
-        nullable=False,
-    )
+    assert result.kind == "Vector"
+    assert result.vector_info == VectorInfo(dim=128)
+    assert result.elem_type == 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.ndarray[tuple[int, ...], np.dtype[np.float32]],
-        attrs=None,
-        nullable=False,
-    )
+    assert result.kind == "Vector"
+    assert result.vector_info == VectorInfo(dim=384)
+    assert result.elem_type == 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:
@@ -105,37 +94,39 @@ def test_ndarray_int64_no_dim() -> None:
     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.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.ndarray[tuple[int, ...], np.dtype[np.float32]],
-        attrs=None,
-        nullable=True,
-    )
+    assert result.kind == "Vector"
+    assert result.vector_info == VectorInfo(dim=None)
+    assert result.elem_type == 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():
+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.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
 
@@ -144,7 +135,7 @@ 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)
 
 
@@ -160,7 +151,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)
 
 
@@ -504,16 +495,16 @@ def test_encode_enriched_type_nullable() -> None:
     assert result["nullable"] is True
 
 
-def test_encode_scalar_numpy_types_schema():
+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 schema["type"]["kind"] == expected_kind, (
+            f"Expected {expected_kind} for {np_type}, got {schema['type']['kind']}"
+        )
         assert not schema.get("nullable", False)