Skip to content

Commit edc6a24

Browse files
lemoragelemorage
committed
refactor: improve closure structure for each type-specific encoder
1 parent e82772b commit edc6a24

File tree

1 file changed

+84
-164
lines changed

1 file changed

+84
-164
lines changed

‎python/cocoindex/convert.py‎

Lines changed: 84 additions & 164 deletions
Original file line numberDiff line numberDiff line change
@@ -82,199 +82,122 @@ def _make_encoder_closure(type_info: AnalyzedTypeInfo | None) -> Callable[[Any],
8282
"""
8383
Create an encoder closure for a specific type.
8484
"""
85-
if type_info is None:
86-
# For untyped encoding, fall back to basic logic
87-
def encode_untyped(value: Any) -> Any:
88-
if dataclasses.is_dataclass(value):
89-
fields = dataclasses.fields(value)
90-
return [
91-
_make_encoder_closure(_get_type_info(f.type))(
92-
getattr(value, f.name)
93-
)
94-
for f in fields
95-
]
96-
97-
if is_namedtuple_type(type(value)):
98-
annotations = type(value).__annotations__
99-
return [
100-
_make_encoder_closure(
101-
_get_type_info(annotations.get(name))
102-
if annotations.get(name)
103-
else None
104-
)(getattr(value, name))
105-
for name in value._fields
106-
]
107-
108-
if isinstance(value, np.number):
109-
return value.item()
110-
111-
if isinstance(value, np.ndarray):
112-
return value
113-
114-
if isinstance(value, (list, tuple)):
115-
return [_make_encoder_closure(None)(v) for v in value]
116-
117-
if isinstance(value, dict):
118-
# Handle empty dict
119-
if not value:
120-
return value
121-
122-
# Handle KTable
123-
first_val = next(iter(value.values()))
124-
if is_struct_type(type(first_val)):
125-
return [
126-
[_make_encoder_closure(None)(k)]
127-
+ _make_encoder_closure(None)(v)
128-
for k, v in value.items()
129-
]
13085

86+
def _encode_numpy(value: Any) -> Any | None:
87+
if isinstance(value, np.number):
88+
return value.item()
89+
if isinstance(value, np.ndarray):
13190
return value
91+
return None
13292

133-
return encode_untyped
134-
135-
variant = type_info.variant
93+
def _encode_struct_like(value: Any) -> Any | None:
94+
if dataclasses.is_dataclass(value):
95+
return [
96+
_make_encoder_closure(_get_type_info(f.type))(getattr(value, f.name))
97+
for f in dataclasses.fields(value)
98+
]
99+
if is_namedtuple_type(type(value)):
100+
annotations = type(value).__annotations__
101+
return [
102+
_make_encoder_closure(
103+
_get_type_info(annotations.get(name))
104+
if annotations.get(name)
105+
else None
106+
)(getattr(value, name))
107+
for name in value._fields
108+
]
109+
return None
136110

137-
# Handle JSON types
138-
if isinstance(variant, AnalyzedBasicType) and variant.kind == "Json":
111+
def _encode_list_or_tuple(value: Any) -> Any | None:
112+
if isinstance(value, (list, tuple)):
113+
encoder = _make_encoder_closure(None)
114+
return [encoder(v) for v in value]
115+
return None
139116

140-
def encode_json_dict(value: Any) -> Any:
141-
if isinstance(value, dict):
142-
# Handle empty dict
143-
if not value:
144-
return value
117+
def _encode_ktable_dict(value: dict[Any, Any]) -> Any | None:
118+
if not value:
119+
return value
120+
first_val = next(iter(value.values()))
121+
if is_struct_type(type(first_val)):
122+
encoder = _make_encoder_closure(None)
123+
return [[encoder(k)] + encoder(v) for k, v in value.items()]
124+
return None
145125

146-
# Handle KTable
147-
first_val = next(iter(value.values()))
148-
if is_struct_type(type(first_val)):
149-
untyped_encoder = _make_encoder_closure(None)
150-
return [
151-
[untyped_encoder(k)] + untyped_encoder(v)
152-
for k, v in value.items()
153-
]
126+
def _encode_untyped(value: Any) -> Any:
127+
if (res := _encode_numpy(value)) is not None:
128+
return res
129+
if (res := _encode_struct_like(value)) is not None:
130+
return res
131+
if (res := _encode_list_or_tuple(value)) is not None:
132+
return res
133+
if isinstance(value, dict):
134+
if (res := _encode_ktable_dict(value)) is not None:
135+
return res
136+
return value
154137

155-
return value
138+
if type_info is None:
139+
return _encode_untyped
156140

157-
return encode_json_dict
141+
variant = type_info.variant
158142

159-
# Handle Any types and special numpy cases
160143
if isinstance(variant, AnalyzedAnyType):
161144

162145
def encode_any_type(value: Any) -> Any:
163-
# Handle numpy types first
164-
if isinstance(value, np.number):
165-
return value.item()
166-
if isinstance(value, np.ndarray):
167-
return value
168-
169-
# Handle tuples - convert to lists for Any type
146+
if (res := _encode_numpy(value)) is not None:
147+
return res
170148
if isinstance(value, tuple):
171149
return [_make_encoder_closure(None)(v) for v in value]
172-
173-
# Handle dataclasses
174-
if dataclasses.is_dataclass(value):
175-
fields = dataclasses.fields(value)
176-
return [
177-
_make_encoder_closure(_get_type_info(f.type))(
178-
getattr(value, f.name)
179-
)
180-
for f in fields
181-
]
182-
183-
# Handle namedtuples
184-
if is_namedtuple_type(type(value)):
185-
annotations = type(value).__annotations__
186-
return [
187-
_make_encoder_closure(
188-
_get_type_info(annotations.get(name))
189-
if annotations.get(name)
190-
else None
191-
)(getattr(value, name))
192-
for name in value._fields
193-
]
194-
195-
# Handle lists
150+
if (res := _encode_struct_like(value)) is not None:
151+
return res
196152
if isinstance(value, list):
197153
return [_make_encoder_closure(None)(v) for v in value]
198-
199-
# Handle dicts
200154
if isinstance(value, dict):
201-
# Handle empty dict
202-
if not value:
203-
return value
204-
205-
# Handle KTable
206-
first_val = next(iter(value.values()))
207-
if is_struct_type(type(first_val)):
208-
return [
209-
[_make_encoder_closure(None)(k)]
210-
+ _make_encoder_closure(None)(v)
211-
for k, v in value.items()
212-
]
213-
155+
if (res := _encode_ktable_dict(value)) is not None:
156+
return res
214157
return value
215158

216159
return encode_any_type
217160

218-
# Handle basic types
219161
if isinstance(variant, AnalyzedBasicType):
220162

221163
def encode_basic_with_numpy(value: Any) -> Any:
222-
# Handle numpy types for basic types
223-
if isinstance(value, np.number):
224-
return value.item()
225-
if isinstance(value, np.ndarray):
226-
return value
164+
if (res := _encode_numpy(value)) is not None:
165+
return res
227166
return value
228167

229168
return encode_basic_with_numpy
230169

231-
# Handle lists
232170
if isinstance(variant, AnalyzedListType):
233-
if variant.elem_type:
234-
elem_encoder = _make_encoder_closure(_get_type_info(variant.elem_type))
235-
return (
236-
lambda value: [elem_encoder(v) for v in value]
237-
if isinstance(value, (list, tuple))
238-
else value
239-
)
240-
else:
241-
fallback_encoder = _make_encoder_closure(None)
242-
return (
243-
lambda value: [fallback_encoder(v) for v in value]
244-
if isinstance(value, (list, tuple))
245-
else value
246-
)
171+
elem_encoder = (
172+
_make_encoder_closure(None)
173+
if not variant.elem_type
174+
else _make_encoder_closure(_get_type_info(variant.elem_type))
175+
)
247176

248-
# Handle dicts
249-
if isinstance(variant, AnalyzedDictType):
250-
if variant.value_type:
251-
value_encoder = _make_encoder_closure(_get_type_info(variant.value_type))
252-
untyped_encoder = _make_encoder_closure(None)
177+
def encode_list(value: Any) -> Any:
178+
if isinstance(value, (list, tuple)):
179+
return [elem_encoder(v) for v in value]
180+
return value
253181

254-
def encode_dict(value: Any) -> Any:
255-
if not isinstance(value, dict):
256-
return value
182+
return encode_list
257183

258-
# Handle empty dict
259-
if not value:
260-
return []
184+
if isinstance(variant, AnalyzedDictType):
185+
if not variant.value_type:
186+
return lambda value: value
261187

262-
# Handle KTable
263-
first_val = next(iter(value.values()))
264-
if is_struct_type(type(first_val)):
265-
return [
266-
[untyped_encoder(k)] + untyped_encoder(v)
267-
for k, v in value.items()
268-
]
188+
value_encoder = _make_encoder_closure(_get_type_info(variant.value_type))
269189

270-
# Handle regular dict
271-
return {k: value_encoder(v) for k, v in value.items()}
190+
def encode_dict(value: Any) -> Any:
191+
if not isinstance(value, dict):
192+
return value
193+
if not value:
194+
return []
195+
if (res := _encode_ktable_dict(value)) is not None:
196+
return res
197+
return {k: value_encoder(v) for k, v in value.items()}
272198

273-
return encode_dict
274-
else:
275-
return lambda value: value
199+
return encode_dict
276200

277-
# Handle struct types
278201
if isinstance(variant, AnalyzedStructType):
279202
struct_type = variant.struct_type
280203

@@ -317,17 +240,14 @@ def encode_namedtuple(value: Any) -> Any:
317240

318241
return encode_namedtuple
319242

320-
# Handle numpy types
321-
def encode_with_numpy_check(value: Any) -> Any:
322-
if isinstance(value, np.number):
323-
return value.item()
324-
if isinstance(value, np.ndarray):
325-
return value
243+
def encode_fallback(value: Any) -> Any:
244+
if (res := _encode_numpy(value)) is not None:
245+
return res
326246
if isinstance(value, tuple):
327247
return [_make_encoder_closure(None)(v) for v in value]
328248
return value
329249

330-
return encode_with_numpy_check
250+
return encode_fallback
331251

332252

333253
def make_engine_value_encoder(type_hint: Type[Any] | str) -> Callable[[Any], Any]:

0 commit comments

Comments
 (0)