Coverage for src / loman / serialization / transformer.py: 97%
372 statements
« prev ^ index » next coverage.py v7.13.5, created at 2026-05-07 21:24 +0000
1"""Object serialization and transformation framework."""
3import contextlib
4import dataclasses
5import graphlib
6import importlib
7from abc import ABC, abstractmethod
8from collections.abc import Callable, Iterable
9from enum import Enum
10from typing import Any
12import numpy as np
13import pandas as pd
15try:
16 import attrs
18 HAS_ATTRS = True
19except ImportError: # pragma: no cover
20 HAS_ATTRS = False
22KEY_TYPE = "type"
23KEY_CLASS = "class"
24KEY_VALUES = "values"
25KEY_DATA = "data"
27TYPENAME_DICT = "dict"
28TYPENAME_TUPLE = "tuple"
29TYPENAME_TRANSFORMABLE = "transformable"
30TYPENAME_ATTRS = "attrs"
31TYPENAME_DATACLASS = "dataclass"
34class UntransformableTypeError(Exception):
35 """Exception raised when a type cannot be transformed for serialization."""
37 pass
40class UnrecognizedTypeError(Exception):
41 """Exception raised when a type is not recognized during transformation."""
43 pass
46class MissingObject:
47 """Sentinel object representing missing or unset values."""
49 def __repr__(self) -> str:
50 """Return string representation of missing object."""
51 return "Missing"
54def order_classes(classes: Iterable[type]) -> list[type]:
55 """Order classes by inheritance hierarchy using topological sort."""
56 graph: dict[type, set[type]] = {x: set() for x in classes}
57 for x in classes:
58 for y in classes:
59 if issubclass(x, y) and x != y:
60 graph[y].add(x)
61 ts = graphlib.TopologicalSorter(graph)
62 return list(ts.static_order())
65class CustomTransformer(ABC):
66 """Abstract base class for custom object transformers."""
68 @property
69 @abstractmethod
70 def name(self) -> str:
71 """Return unique name identifier for this transformer."""
72 pass # pragma: no cover
74 @abstractmethod
75 def to_dict(self, transformer: "Transformer", o: object) -> dict[str, Any]:
76 """Convert object to dictionary representation."""
77 pass # pragma: no cover
79 @abstractmethod
80 def from_dict(self, transformer: "Transformer", d: dict[str, Any]) -> object:
81 """Reconstruct object from dictionary representation."""
82 pass # pragma: no cover
84 @property
85 def supported_direct_types(self) -> Iterable[type]:
86 """Return types that this transformer handles directly."""
87 return []
89 @property
90 def supported_subtypes(self) -> Iterable[Any]:
91 """Return base types whose subtypes this transformer can handle."""
92 return []
95class Transformable(ABC):
96 """Abstract base class for objects that can transform themselves."""
98 @abstractmethod
99 def to_dict(self, transformer: "Transformer") -> dict[str, Any]:
100 """Convert this object to dictionary representation."""
101 pass # pragma: no cover
103 @classmethod
104 @abstractmethod
105 def from_dict(cls, transformer: "Transformer", d: dict[str, Any]) -> object:
106 """Reconstruct object from dictionary representation."""
107 pass # pragma: no cover
110class Transformer:
111 """Main transformer class for object serialization and deserialization."""
113 def __init__(self, *, strict: bool = True) -> None:
114 """Initialize transformer with strict mode setting."""
115 self.strict = strict
117 self._direct_type_map: dict[type, CustomTransformer] = {}
118 self._subtype_order: list[type] = []
119 self._subtype_map: dict[type, CustomTransformer] = {}
120 self._transformers: dict[str, CustomTransformer] = {}
121 self._transformable_types: dict[str, type[Transformable]] = {}
122 self._attrs_types: dict[str, type] = {}
123 self._dataclass_types: dict[str, type] = {}
125 def register(self, t: CustomTransformer | type[Transformable] | type) -> None:
126 """Register a transformer, transformable type, or regular type."""
127 if isinstance(t, CustomTransformer):
128 self.register_transformer(t)
129 elif isinstance(t, type) and issubclass(t, Transformable):
130 self.register_transformable(t)
131 elif HAS_ATTRS and isinstance(t, type) and attrs.has(t):
132 self.register_attrs(t)
133 elif isinstance(t, type) and dataclasses.is_dataclass(t):
134 self.register_dataclass(t)
135 else:
136 msg = f"Unable to register {t}"
137 raise ValueError(msg)
139 def register_transformer(self, transformer: CustomTransformer) -> None:
140 """Register a custom transformer for specific types."""
141 assert transformer.name not in self._transformers # noqa: S101
142 for type_ in transformer.supported_direct_types:
143 assert type_ not in self._direct_type_map # noqa: S101
144 for type_ in transformer.supported_subtypes:
145 assert type_ not in self._subtype_map # noqa: S101
147 self._transformers[transformer.name] = transformer
149 for type_ in transformer.supported_direct_types:
150 self._direct_type_map[type_] = transformer
152 contains_supported_subtypes = False
153 for type_ in transformer.supported_subtypes:
154 contains_supported_subtypes = True
155 self._subtype_map[type_] = transformer
156 if contains_supported_subtypes:
157 self._subtype_order = order_classes(self._subtype_map.keys())
159 def register_transformable(self, transformable_type: type[Transformable]) -> None:
160 """Register a transformable type that can serialize itself."""
161 name = transformable_type.__name__
162 assert name not in self._transformable_types # noqa: S101
163 self._transformable_types[name] = transformable_type
165 def register_attrs(self, attrs_type: type) -> None:
166 """Register an attrs-decorated class for serialization."""
167 name = attrs_type.__name__
168 assert name not in self._attrs_types # noqa: S101
169 self._attrs_types[name] = attrs_type
171 def register_dataclass(self, dataclass_type: type) -> None:
172 """Register a dataclass for serialization."""
173 name = dataclass_type.__name__
174 assert name not in self._dataclass_types # noqa: S101
175 self._dataclass_types[name] = dataclass_type
177 def get_transformer_for_obj(self, obj: object) -> CustomTransformer | None:
178 """Get the appropriate transformer for a given object."""
179 transformer = self._direct_type_map.get(type(obj))
180 if transformer is not None:
181 return transformer
182 for tp in self._subtype_order:
183 if isinstance(obj, tp):
184 return self._subtype_map[tp]
185 return None
187 def get_transformer_for_name(self, name: str) -> CustomTransformer | None:
188 """Get a transformer by its registered name."""
189 transformer = self._transformers.get(name)
190 return transformer
192 def to_dict(self, o: object) -> Any:
193 """Convert an object to a serializable dictionary representation."""
194 if isinstance(o, str) or o is None or o is True or o is False or isinstance(o, (int, float)):
195 return o
196 elif isinstance(o, tuple):
197 return {KEY_TYPE: TYPENAME_TUPLE, KEY_VALUES: [self.to_dict(x) for x in o]}
198 elif isinstance(o, list):
199 return [self.to_dict(x) for x in o]
200 elif isinstance(o, dict):
201 return self._dict_to_dict(o)
202 # Check registered custom transformers before generic dataclass/attrs paths
203 # so that explicitly registered types (e.g. NodeKey) take priority.
204 elif self.get_transformer_for_obj(o) is not None:
205 return self._to_dict_transformer(o)
206 elif isinstance(o, Transformable):
207 return {KEY_TYPE: TYPENAME_TRANSFORMABLE, KEY_CLASS: type(o).__name__, KEY_DATA: o.to_dict(self)}
208 elif HAS_ATTRS and attrs.has(type(o)):
209 return self._attrs_to_dict(o)
210 elif dataclasses.is_dataclass(o) and not isinstance(o, type):
211 return self._dataclass_to_dict(o)
212 else:
213 return self._to_dict_transformer(o)
215 def _dict_to_dict(self, o: dict[Any, Any]) -> dict[str, Any]:
216 """Convert a dictionary to serializable form."""
217 d = {k: self.to_dict(v) for k, v in o.items()}
218 if KEY_TYPE in o:
219 return {KEY_TYPE: TYPENAME_DICT, KEY_DATA: d}
220 else:
221 return d
223 def _attrs_to_dict(self, o: object) -> dict[str, Any]:
224 """Convert an attrs object to serializable dictionary form."""
225 data: dict[str, Any] = {}
226 for a in o.__attrs_attrs__: # type: ignore[attr-defined]
227 data[a.name] = self.to_dict(o.__getattribute__(a.name))
228 res: dict[str, Any] = {KEY_TYPE: TYPENAME_ATTRS, KEY_CLASS: type(o).__name__}
229 if len(data) > 0:
230 res[KEY_DATA] = data
231 return res
233 def _dataclass_to_dict(self, o: object) -> dict[str, Any]:
234 """Convert a dataclass object to serializable dictionary form."""
235 data: dict[str, Any] = {}
236 for f in dataclasses.fields(o): # type: ignore[arg-type]
237 data[f.name] = self.to_dict(getattr(o, f.name))
238 res: dict[str, Any] = {KEY_TYPE: TYPENAME_DATACLASS, KEY_CLASS: type(o).__name__}
239 if len(data) > 0:
240 res[KEY_DATA] = data
241 return res
243 def _to_dict_transformer(self, o: object) -> dict[str, Any] | None:
244 """Convert an object using a registered custom transformer."""
245 transformer = self.get_transformer_for_obj(o)
246 if transformer is None:
247 if self.strict:
248 msg = f"Could not transform object of type {type(o).__name__}"
249 raise UntransformableTypeError(msg)
250 else:
251 return None
252 d = transformer.to_dict(self, o)
253 d[KEY_TYPE] = transformer.name
254 return d
256 def from_dict(self, d: Any) -> Any:
257 """Convert a dictionary representation back to the original object."""
258 if isinstance(d, str) or d is None or d is True or d is False or isinstance(d, (int, float)):
259 return d
260 elif isinstance(d, list):
261 return [self.from_dict(x) for x in d]
262 elif isinstance(d, dict):
263 type_ = d.get(KEY_TYPE)
264 if type_ is None:
265 return {k: self.from_dict(v) for k, v in d.items()}
266 elif type_ == TYPENAME_TUPLE:
267 return tuple(self.from_dict(x) for x in d[KEY_VALUES])
268 elif type_ == TYPENAME_DICT:
269 return {k: self.from_dict(v) for k, v in d[KEY_DATA].items()}
270 elif type_ == TYPENAME_TRANSFORMABLE:
271 return self._from_dict_transformable(d)
272 elif type_ == TYPENAME_ATTRS:
273 return self._from_attrs(d)
274 elif type_ == TYPENAME_DATACLASS:
275 return self._from_dataclass(d)
276 else:
277 return self._from_dict_transformer(type_, d)
278 else:
279 msg = "Unable to determine object type from dictionary"
280 raise ValueError(msg)
282 def _from_dict_transformable(self, d: dict[str, Any]) -> object:
283 """Reconstruct a Transformable object from dictionary form."""
284 classname = d[KEY_CLASS]
285 cls = self._transformable_types.get(classname)
286 if cls is None:
287 if self.strict:
288 msg = f"Unable to transform Transformable object of class {classname}"
289 raise UnrecognizedTypeError(msg)
290 else:
291 return MissingObject()
292 else:
293 return cls.from_dict(self, d[KEY_DATA])
295 def _from_attrs(self, d: dict[str, Any]) -> object:
296 """Reconstruct an attrs object from dictionary form."""
297 if not HAS_ATTRS: # pragma: no cover
298 if self.strict:
299 msg = "attrs package not installed"
300 raise UnrecognizedTypeError(msg)
301 return MissingObject()
302 cls = self._attrs_types.get(d[KEY_CLASS])
303 if cls is None:
304 if self.strict:
305 msg = f"Unable to create attrs object of type {cls}"
306 raise UnrecognizedTypeError(msg)
307 else:
308 return MissingObject()
309 else:
310 kwargs: dict[str, Any] = {}
311 if KEY_DATA in d:
312 for key, value in d[KEY_DATA].items():
313 kwargs[key] = self.from_dict(value)
314 return cls(**kwargs)
316 def _from_dataclass(self, d: dict[str, Any]) -> object:
317 """Reconstruct a dataclass object from dictionary form."""
318 cls = self._dataclass_types.get(d[KEY_CLASS])
319 if cls is None:
320 if self.strict:
321 msg = f"Unable to create dataclass object of type {cls}"
322 raise UnrecognizedTypeError(msg)
323 else:
324 return MissingObject()
325 else:
326 kwargs: dict[str, Any] = {}
327 if KEY_DATA in d:
328 for key, value in d[KEY_DATA].items():
329 kwargs[key] = self.from_dict(value)
330 return cls(**kwargs)
332 def _from_dict_transformer(self, type_: str, d: dict[str, Any]) -> object:
333 """Reconstruct an object using a registered custom transformer."""
334 transformer = self.get_transformer_for_name(type_)
335 if transformer is None:
336 if self.strict:
337 msg = f"Unable to transform object of type {type_}"
338 raise UnrecognizedTypeError(msg)
339 else:
340 return MissingObject()
341 return transformer.from_dict(self, d)
344class NdArrayTransformer(CustomTransformer):
345 """Transformer for NumPy ndarray objects."""
347 @property
348 def name(self) -> str:
349 """Return transformer name."""
350 return "ndarray"
352 def to_dict(self, transformer: "Transformer", o: object) -> dict[str, Any]:
353 """Convert numpy array to dictionary with shape, dtype, and data."""
354 assert isinstance(o, np.ndarray) # noqa: S101
355 return {"shape": list(o.shape), "dtype": o.dtype.str, "data": transformer.to_dict(o.ravel().tolist())} # type: ignore[arg-type]
357 def from_dict(self, transformer: "Transformer", d: dict[str, Any]) -> object:
358 """Reconstruct numpy array from dictionary."""
359 return np.array(transformer.from_dict(d["data"]), d["dtype"]).reshape(d["shape"])
361 @property
362 def supported_direct_types(self) -> Iterable[type]:
363 """Return supported numpy array types."""
364 return [np.ndarray]
367class EnumTransformer(CustomTransformer):
368 """Transformer for Enum subclasses.
370 Enum classes must be registered via :meth:`register_enum` before use.
371 """
373 def __init__(self) -> None:
374 """Initialise with an empty enum registry."""
375 self._registry: dict[str, type[Enum]] = {}
377 def register_enum(self, enum_class: type[Enum]) -> None:
378 """Register an enum class so its members can be deserialized."""
379 self._registry[enum_class.__qualname__] = enum_class
381 @property
382 def name(self) -> str:
383 """Return transformer name."""
384 return "enum"
386 def to_dict(self, transformer: "Transformer", o: object) -> dict[str, Any]:
387 """Convert an Enum member to a dict with class qualname and member name."""
388 assert isinstance(o, Enum) # noqa: S101
389 return {"enum_class": type(o).__qualname__, "value": o.name}
391 def from_dict(self, transformer: "Transformer", d: dict[str, Any]) -> object:
392 """Reconstruct an Enum member from its serialized form."""
393 enum_class = self._registry.get(d["enum_class"])
394 if enum_class is None:
395 msg = f"Unknown enum class: {d['enum_class']!r}. Register it with EnumTransformer.register_enum()."
396 raise UnrecognizedTypeError(msg)
397 return enum_class[d["value"]]
399 @property
400 def supported_subtypes(self) -> Iterable[type]:
401 """Handle all Enum subclasses."""
402 return [Enum]
405class FunctionRefTransformer(CustomTransformer):
406 """Transformer for importable callables (module-level functions and methods).
408 Lambdas and closures (whose ``__qualname__`` contains ``<lambda>`` or
409 ``<locals>``) are explicitly rejected with a :class:`ValueError`.
410 """
412 @property
413 def name(self) -> str:
414 """Return transformer name."""
415 return "func_ref"
417 def to_dict(self, transformer: "Transformer", o: object) -> dict[str, Any]:
418 """Serialize a callable as its module path and qualname."""
419 if not callable(o):
420 msg = f"Object {o!r} is not callable"
421 raise TypeError(msg)
422 qualname = getattr(o, "__qualname__", None)
423 module = getattr(o, "__module__", None)
424 if qualname is None or module is None:
425 msg = f"Cannot serialize {o!r}: missing __qualname__ or __module__"
426 raise ValueError(msg)
427 if "<lambda>" in qualname:
428 msg = f"Cannot serialize lambda function {o!r}: lambdas are not importable"
429 raise ValueError(msg)
430 if "<locals>" in qualname:
431 msg = f"Cannot serialize closure/local function {o!r}: non-importable"
432 raise ValueError(msg)
433 # Verify the callable is actually reachable via import before committing.
434 try:
435 mod = importlib.import_module(module)
436 obj: Any = mod
437 for part in qualname.split("."):
438 obj = getattr(obj, part)
439 if obj is not o:
440 msg = f"Cannot serialize {o!r}: import round-trip returned a different object"
441 raise ValueError(msg)
442 except (ImportError, AttributeError) as exc:
443 msg = f"Cannot serialize {o!r}: not importable ({exc})"
444 raise ValueError(msg) from exc
445 return {"module": module, "qualname": qualname}
447 def from_dict(self, transformer: "Transformer", d: dict[str, Any]) -> object:
448 """Reconstruct a callable from its module path and qualname."""
449 module = importlib.import_module(d["module"])
450 obj: Any = module
451 for part in d["qualname"].split("."):
452 obj = getattr(obj, part)
453 return obj
455 @property
456 def supported_direct_types(self) -> Iterable[type]:
457 """Register the built-in function types explicitly handled."""
458 # We use supported_subtypes for the broad callable match instead,
459 # but we must list at least one concrete type here to help dispatch.
460 # The broad subtype match on Callable covers everything callable.
461 return []
463 @property
464 def supported_subtypes(self) -> Iterable[Any]:
465 """Match all callables via Callable ABC."""
466 return [Callable]
469class DillFunctionTransformer(CustomTransformer):
470 """Transformer that serializes any callable — including lambdas and closures — using dill.
472 The callable is serialized with :func:`dill.dumps` and the resulting bytes
473 are stored as a base64-encoded string inside the JSON document. On load the
474 bytes are decoded and passed to :func:`dill.loads`.
476 .. note::
477 The embedded dill blob is **not** portable across Python versions and
478 shares the same stability caveats as :meth:`~loman.Computation.write_dill`.
479 Register this transformer when convenient lambda/closure round-trips matter
480 more than portability.
482 Example::
484 from loman import Computation, ComputationSerializer
485 from loman.serialization import DillFunctionTransformer
487 s = ComputationSerializer(use_dill_for_functions=True)
488 comp = Computation()
489 comp.add_node('a', value=1)
490 comp.add_node('b', lambda a: a + 1)
491 comp.compute_all()
492 comp.write_json('comp.json', serializer=s)
493 comp2 = Computation.read_json('comp.json', serializer=s)
494 assert comp2.v.b == 2
495 """
497 @property
498 def name(self) -> str:
499 """Return transformer name."""
500 return "dill_func"
502 def to_dict(self, transformer: "Transformer", o: object) -> dict[str, Any]:
503 """Serialize a callable to a base64-encoded dill blob."""
504 import base64
506 import dill # nosec B403 # dill is a trusted dependency for this specific use case and most likely be deprecated in the future in favor of a more portable solution, so we allow it here with a blanket nosec directive
508 if not callable(o):
509 msg = f"Object {o!r} is not callable"
510 raise TypeError(msg)
511 blob = dill.dumps(o)
512 return {"blob": base64.b64encode(blob).decode("ascii")}
514 def from_dict(self, transformer: "Transformer", d: dict[str, Any]) -> object:
515 """Reconstruct a callable from a base64-encoded dill blob."""
516 import base64
518 import dill # nosec B403 # dill is a trusted dependency for this specific use case and most likely be deprecated in the future in favor of a more portable solution, so we allow it here with a blanket nosec directive
520 blob = base64.b64decode(d["blob"].encode("ascii"))
521 return dill.loads(blob) # noqa: S301 # nosec B301
523 @property
524 def supported_direct_types(self) -> Iterable[type]:
525 """No direct type matches — rely on subtype matching."""
526 return []
528 @property
529 def supported_subtypes(self) -> Iterable[Any]:
530 """Match all callables via Callable ABC."""
531 return [Callable]
534class DataFrameTransformer(CustomTransformer):
535 """Transformer for :class:`pandas.DataFrame` objects."""
537 @property
538 def name(self) -> str:
539 """Return transformer name."""
540 return "dataframe"
542 def to_dict(self, transformer: "Transformer", o: object) -> dict[str, Any]:
543 """Serialize a DataFrame using split orientation."""
544 assert isinstance(o, pd.DataFrame) # noqa: S101
545 return {
546 "columns": list(o.columns),
547 "index": transformer.to_dict(list(o.index)),
548 "data": transformer.to_dict(o.values.tolist()),
549 "dtypes": {col: str(dtype) for col, dtype in o.dtypes.items()},
550 }
552 def from_dict(self, transformer: "Transformer", d: dict[str, Any]) -> object:
553 """Reconstruct a DataFrame from its serialized form."""
554 data = transformer.from_dict(d["data"])
555 columns = d["columns"]
556 index = transformer.from_dict(d["index"])
557 dtypes = d.get("dtypes", {})
558 df = pd.DataFrame(data, columns=columns, index=index)
559 for col, dtype in dtypes.items():
560 with contextlib.suppress(ValueError, TypeError): # pragma: no cover
561 df[col] = df[col].astype(dtype)
562 return df
564 @property
565 def supported_direct_types(self) -> Iterable[type]:
566 """Return supported pandas DataFrame type."""
567 return [pd.DataFrame]
570class SeriesTransformer(CustomTransformer):
571 """Transformer for :class:`pandas.Series` objects."""
573 @property
574 def name(self) -> str:
575 """Return transformer name."""
576 return "series"
578 def to_dict(self, transformer: "Transformer", o: object) -> dict[str, Any]:
579 """Serialize a Series with its name, dtype, index, and data."""
580 assert isinstance(o, pd.Series) # noqa: S101
581 return {
582 "name": o.name,
583 "dtype": str(o.dtype),
584 "index": transformer.to_dict(list(o.index)),
585 "data": transformer.to_dict(o.tolist()),
586 }
588 def from_dict(self, transformer: "Transformer", d: dict[str, Any]) -> object:
589 """Reconstruct a Series from its serialized form."""
590 data = transformer.from_dict(d["data"])
591 index = transformer.from_dict(d["index"])
592 s = pd.Series(data, index=index, name=d.get("name"))
593 with contextlib.suppress(ValueError, TypeError): # pragma: no cover
594 s = s.astype(d["dtype"])
595 return s
597 @property
598 def supported_direct_types(self) -> Iterable[type]:
599 """Return supported pandas Series type."""
600 return [pd.Series]
603class NodeKeyTransformer(CustomTransformer):
604 """Transformer for :class:`~loman.nodekey.NodeKey` objects."""
606 @property
607 def name(self) -> str:
608 """Return transformer name."""
609 return "nodekey"
611 def to_dict(self, transformer: "Transformer", o: object) -> dict[str, Any]:
612 """Serialize a NodeKey as its path string."""
613 return {"path": str(o)}
615 def from_dict(self, transformer: "Transformer", d: dict[str, Any]) -> object:
616 """Reconstruct a NodeKey from its path string."""
617 from loman.nodekey import parse_nodekey
619 return parse_nodekey(d["path"])
621 @property
622 def supported_direct_types(self) -> Iterable[type]:
623 """Return supported NodeKey type."""
624 from loman.nodekey import NodeKey
626 return [NodeKey]