Skip to content

Commit 0994ac7

Browse files
committed
Code cleanups for basic type encoding in Python SDK.
1 parent 4ddb948 commit 0994ac7

File tree

1 file changed

+25
-26
lines changed

1 file changed

+25
-26
lines changed

python/cocoindex/typing.py

Lines changed: 25 additions & 26 deletions
Original file line numberDiff line numberDiff line change
@@ -5,9 +5,8 @@
55
class Vector(NamedTuple):
66
dim: int | None
77

8-
class NumBits(NamedTuple):
9-
bits: int
10-
8+
class TypeKind(NamedTuple):
9+
kind: str
1110
class TypeAttr:
1211
key: str
1312
value: Any
@@ -16,10 +15,10 @@ def __init__(self, key: str, value: Any):
1615
self.key = key
1716
self.value = value
1817

19-
Float32 = Annotated[float, NumBits(32)]
20-
Float64 = Annotated[float, NumBits(64)]
21-
Range = Annotated[tuple[int, int], 'range']
22-
Json = Annotated[Any, 'json']
18+
Float32 = Annotated[float, TypeKind('Float32')]
19+
Float64 = Annotated[float, TypeKind('Float64')]
20+
Range = Annotated[tuple[int, int], TypeKind('Range')]
21+
Json = Annotated[Any, TypeKind('Json')]
2322

2423
def _find_annotation(metadata, cls):
2524
for m in iter(metadata):
@@ -32,36 +31,36 @@ def _get_origin_type_and_metadata(t):
3231
return (t.__origin__, t.__metadata__)
3332
return (t, ())
3433

35-
def _basic_type_to_json_value(t, metadata):
34+
def _type_to_json_value(t, metadata):
3635
origin_type = typing.get_origin(t)
3736
if origin_type is collections.abc.Sequence or origin_type is list:
3837
dim = _find_annotation(metadata, Vector)
3938
if dim is None:
4039
raise ValueError(f"Vector dimension not found for {t}")
4140
args = typing.get_args(t)
41+
origin_type, metadata = _get_origin_type_and_metadata(args[0])
4242
type_json = {
4343
'kind': 'Vector',
44-
'element_type': _basic_type_to_json_value(*_get_origin_type_and_metadata(args[0])),
44+
'element_type': _type_to_json_value(origin_type, metadata),
4545
'dimension': dim.dim,
4646
}
4747
else:
48-
if t is bytes:
49-
kind = 'Bytes'
50-
elif t is str:
51-
kind = 'Str'
52-
elif t is bool:
53-
kind = 'Bool'
54-
elif t is int:
55-
kind = 'Int64'
56-
elif t is float:
57-
num_bits = _find_annotation(metadata, NumBits)
58-
kind = 'Float32' if num_bits is not None and num_bits.bits <= 32 else 'Float64'
59-
elif t is Range:
60-
kind = 'Range'
61-
elif t is Json:
62-
kind = 'Json'
48+
type_kind = _find_annotation(metadata, TypeKind)
49+
if type_kind is not None:
50+
kind = type_kind.kind
6351
else:
64-
raise ValueError(f"type unsupported yet: {t}")
52+
if t is bytes:
53+
kind = 'Bytes'
54+
elif t is str:
55+
kind = 'Str'
56+
elif t is bool:
57+
kind = 'Bool'
58+
elif t is int:
59+
kind = 'Int64'
60+
elif t is float:
61+
kind = 'Float64'
62+
else:
63+
raise ValueError(f"type unsupported yet: {t}")
6564
type_json = { 'kind': kind }
6665

6766
return type_json
@@ -70,7 +69,7 @@ def _enriched_type_to_json_value(t) -> dict[str, Any] | None:
7069
if t is None:
7170
return None
7271
t, metadata = _get_origin_type_and_metadata(t)
73-
enriched_type_json = {'type': _basic_type_to_json_value(t, metadata)}
72+
enriched_type_json = {'type': _type_to_json_value(t, metadata)}
7473
attrs = None
7574
for attr in metadata:
7675
if isinstance(attr, TypeAttr):

0 commit comments

Comments
 (0)