diff --git a/python/cocoindex/convert.py b/python/cocoindex/convert.py index 1a254ac2..946153e0 100644 --- a/python/cocoindex/convert.py +++ b/python/cocoindex/convert.py @@ -92,10 +92,14 @@ def make_engine_value_decoder( if src_type_kind == "Struct": return _make_engine_struct_to_dict_decoder(field_path, src_type["fields"]) if src_type_kind in TABLE_TYPES: - raise ValueError( - f"Missing type annotation for `{''.join(field_path)}`." - f"It's required for {src_type_kind} type." - ) + if src_type_kind == "LTable": + return _make_engine_ltable_to_list_dict_decoder( + field_path, src_type["row"]["fields"] + ) + elif src_type_kind == "KTable": + return _make_engine_ktable_to_dict_dict_decoder( + field_path, src_type["row"]["fields"] + ) return lambda value: value # Handle struct -> dict binding for explicit dict annotations @@ -340,6 +344,77 @@ def decode_to_dict(values: list[Any] | None) -> dict[str, Any] | None: return decode_to_dict +def _make_engine_ltable_to_list_dict_decoder( + field_path: list[str], + src_fields: list[dict[str, Any]], +) -> Callable[[list[Any] | None], list[dict[str, Any]] | None]: + """Make a decoder from engine LTable values to a list of dicts.""" + + # Create a decoder for each row (struct) to dict + row_decoder = _make_engine_struct_to_dict_decoder(field_path, src_fields) + + def decode_to_list_dict(values: list[Any] | None) -> list[dict[str, Any]] | None: + if values is None: + return None + result = [] + for i, row_values in enumerate(values): + decoded_row = row_decoder(row_values) + if decoded_row is None: + raise ValueError( + f"LTable row at index {i} decoded to None, which is not allowed." + ) + result.append(decoded_row) + return result + + return decode_to_list_dict + + +def _make_engine_ktable_to_dict_dict_decoder( + field_path: list[str], + src_fields: list[dict[str, Any]], +) -> Callable[[list[Any] | None], dict[Any, dict[str, Any]] | None]: + """Make a decoder from engine KTable values to a dict of dicts.""" + + if not src_fields: + raise ValueError("KTable must have at least one field for the key") + + # First field is the key, remaining fields are the value + key_field_schema = src_fields[0] + value_fields_schema = src_fields[1:] + + # Create decoders + field_path.append(f".{key_field_schema.get('name', KEY_FIELD_NAME)}") + key_decoder = make_engine_value_decoder(field_path, key_field_schema["type"], Any) + field_path.pop() + + value_decoder = _make_engine_struct_to_dict_decoder(field_path, value_fields_schema) + + def decode_to_dict_dict( + values: list[Any] | None, + ) -> dict[Any, dict[str, Any]] | None: + if values is None: + return None + result = {} + for row_values in values: + if not row_values: + raise ValueError("KTable row must have at least 1 value (the key)") + key = key_decoder(row_values[0]) + if len(row_values) == 1: + value: dict[str, Any] = {} + else: + tmp = value_decoder(row_values[1:]) + if tmp is None: + value = {} + else: + value = tmp + if isinstance(key, dict): + key = tuple(key.values()) + result[key] = value + return result + + return decode_to_dict_dict + + def dump_engine_object(v: Any) -> Any: """Recursively dump an object for engine. Engine side uses `Pythonized` to catch.""" if v is None: diff --git a/python/cocoindex/tests/test_convert.py b/python/cocoindex/tests/test_convert.py index 622b5c88..57260a55 100644 --- a/python/cocoindex/tests/test_convert.py +++ b/python/cocoindex/tests/test_convert.py @@ -1341,3 +1341,130 @@ class Point(NamedTuple): validate_full_roundtrip( instance, Point, (expected_dict, dict), (expected_dict, Any) ) + + +def test_roundtrip_ltable_to_list_dict_binding() -> None: + """Test LTable -> list[dict] binding with Any annotation.""" + + @dataclass + class User: + id: str + name: str + age: int + + users = [User("u1", "Alice", 25), User("u2", "Bob", 30), User("u3", "Charlie", 35)] + expected_list_dict = [ + {"id": "u1", "name": "Alice", "age": 25}, + {"id": "u2", "name": "Bob", "age": 30}, + {"id": "u3", "name": "Charlie", "age": 35}, + ] + + # Test Any annotation + validate_full_roundtrip(users, list[User], (expected_list_dict, Any)) + + +def test_roundtrip_ktable_to_dict_dict_binding() -> None: + """Test KTable -> dict[K, dict] binding with Any annotation.""" + + @dataclass + class Product: + name: str + price: float + active: bool + + products = { + "p1": Product("Widget", 29.99, True), + "p2": Product("Gadget", 49.99, False), + "p3": Product("Tool", 19.99, True), + } + expected_dict_dict = { + "p1": {"name": "Widget", "price": 29.99, "active": True}, + "p2": {"name": "Gadget", "price": 49.99, "active": False}, + "p3": {"name": "Tool", "price": 19.99, "active": True}, + } + + # Test Any annotation + validate_full_roundtrip(products, dict[str, Product], (expected_dict_dict, Any)) + + +def test_roundtrip_ktable_with_complex_key() -> None: + """Test KTable with complex key types -> dict binding.""" + + @dataclass(frozen=True) + class OrderKey: + shop_id: str + version: int + + @dataclass + class Order: + customer: str + total: float + + orders = { + OrderKey("shop1", 1): Order("Alice", 100.0), + OrderKey("shop2", 2): Order("Bob", 200.0), + } + expected_dict_dict = { + ("shop1", 1): {"customer": "Alice", "total": 100.0}, + ("shop2", 2): {"customer": "Bob", "total": 200.0}, + } + + # Test Any annotation + validate_full_roundtrip(orders, dict[OrderKey, Order], (expected_dict_dict, Any)) + + +def test_roundtrip_ltable_with_nested_structs() -> None: + """Test LTable with nested structs -> list[dict] binding.""" + + @dataclass + class Address: + street: str + city: str + + @dataclass + class Person: + name: str + age: int + address: Address + + people = [ + Person("John", 30, Address("123 Main St", "Anytown")), + Person("Jane", 25, Address("456 Oak Ave", "Somewhere")), + ] + expected_list_dict = [ + { + "name": "John", + "age": 30, + "address": {"street": "123 Main St", "city": "Anytown"}, + }, + { + "name": "Jane", + "age": 25, + "address": {"street": "456 Oak Ave", "city": "Somewhere"}, + }, + ] + + # Test Any annotation + validate_full_roundtrip(people, list[Person], (expected_list_dict, Any)) + + +def test_roundtrip_ktable_with_list_fields() -> None: + """Test KTable with list fields -> dict binding.""" + + @dataclass + class Team: + name: str + members: list[str] + active: bool + + teams = { + "team1": Team("Dev Team", ["Alice", "Bob"], True), + "team2": Team("QA Team", ["Charlie", "David"], False), + } + expected_dict_dict = { + "team1": {"name": "Dev Team", "members": ["Alice", "Bob"], "active": True}, + "team2": {"name": "QA Team", "members": ["Charlie", "David"], "active": False}, + } + + # Test Any annotation + validate_full_roundtrip(teams, dict[str, Team], (expected_dict_dict, Any))