Coverage for src/loman/serialization/transformer.py: 94%
599 statements
« prev ^ index » next coverage.py v7.16.0, created at 2026-09-07 00:32 +0000
1"""Object serialization and transformation framework."""
3import contextlib
4import contextvars
5import dataclasses
6import graphlib
7import importlib
8import math
9import types
10from abc import ABC, abstractmethod
11from collections.abc import Callable, Iterable, Iterator
12from enum import Enum
13from typing import Any, NamedTuple, cast
15import numpy as np
16import pandas as pd
18try:
19 import attrs
21 HAS_ATTRS = True
22except ImportError: # pragma: no cover
23 HAS_ATTRS = False
25KEY_TYPE = "type"
26KEY_CLASS = "class"
27KEY_VALUES = "values"
28KEY_DATA = "data"
29KEY_ITEMS = "items"
31TYPENAME_DICT = "dict"
32TYPENAME_TUPLE = "tuple"
33TYPENAME_TRANSFORMABLE = "transformable"
34TYPENAME_ATTRS = "attrs"
35TYPENAME_DATACLASS = "dataclass"
36TYPENAME_FLOAT = "float"
38# Keys that carry meaning in an encoded value. A user dict containing any of
39# them is written in the escaped ``{"type": "dict", ...}`` form so that its
40# contents can never be mistaken for encoding metadata on the way back in.
41# "$blob" is here for the same reason "type" is: a dict that happens to contain
42# it would otherwise be read back as a reference to someone else's bytes.
43RESERVED_KEYS = frozenset({KEY_TYPE, "$blob"})
45# Non-finite floats have no JSON literal. Writing them as bare NaN / Infinity
46# tokens --- json.dump's default --- produces a file Python can read back and
47# nothing else can, so they are tagged and ``allow_nan=False`` is set, which
48# makes an invalid document structurally impossible rather than merely unlikely.
49_NONFINITE_TO_NAME = {"nan": "NaN", "inf": "Infinity", "-inf": "-Infinity"}
50_NAME_TO_FLOAT = {"NaN": float("nan"), "Infinity": float("inf"), "-Infinity": float("-inf")}
53class _WriteScope(NamedTuple):
54 """Ambient state for one encoding pass: where blobs go, and for which node."""
56 sink: Any
57 profile: Any
58 node: str | None
59 tags: frozenset[str]
60 store: str | None
63# Set by Transformer.writing() / .reading(). ContextVars rather than instance
64# attributes so that two threads encoding through the same Transformer cannot
65# see each other's sink.
66_WRITE_SCOPE: contextvars.ContextVar[_WriteScope | None] = contextvars.ContextVar("loman_blob_write", default=None)
67_READ_SCOPE: contextvars.ContextVar[Any] = contextvars.ContextVar("loman_blob_read", default=None)
70class UntransformableTypeError(Exception):
71 """Exception raised when a type cannot be transformed for serialization."""
74class UnrecognizedTypeError(Exception):
75 """Exception raised when a type is not recognized during transformation."""
78class DuplicateRegistrationError(ValueError):
79 """Exception raised when a name or type is registered on a Transformer twice.
81 Subclasses :class:`ValueError` because that is what :meth:`Transformer.register`
82 already raises for a type it cannot register at all.
83 """
86class MissingObject:
87 """Sentinel object representing missing or unset values."""
89 def __repr__(self) -> str:
90 """Return string representation of missing object."""
91 return "Missing"
94def _nonfinite_name(value: float) -> str:
95 """Return the JSON-safe name for a non-finite float."""
96 return _NONFINITE_TO_NAME[repr(value)]
99def order_classes(classes: Iterable[type]) -> list[type]:
100 """Order classes by inheritance hierarchy using topological sort."""
101 graph: dict[type, set[type]] = {x: set() for x in classes}
102 for x in classes:
103 for y in classes:
104 if issubclass(x, y) and x != y:
105 graph[y].add(x)
106 ts = graphlib.TopologicalSorter(graph)
107 return list(ts.static_order())
110class CustomTransformer(ABC):
111 """Abstract base class for custom object transformers."""
113 @property
114 @abstractmethod
115 def name(self) -> str:
116 """Return unique name identifier for this transformer."""
117 # pragma: no cover
119 @abstractmethod
120 def to_dict(self, transformer: "Transformer", o: object) -> dict[str, Any]:
121 """Convert object to dictionary representation."""
122 # pragma: no cover
124 @abstractmethod
125 def from_dict(self, transformer: "Transformer", d: dict[str, Any]) -> object:
126 """Reconstruct object from dictionary representation."""
127 # pragma: no cover
129 @property
130 def supported_direct_types(self) -> Iterable[type]:
131 """Return types that this transformer handles directly."""
132 return []
134 @property
135 def supported_subtypes(self) -> Iterable[Any]:
136 """Return base types whose subtypes this transformer can handle."""
137 return []
140class SimpleTransformer(CustomTransformer):
141 """A :class:`CustomTransformer` built from two plain functions.
143 Implementing :class:`CustomTransformer` directly means a subclass with three
144 members. When the encoding is just "turn it into a dict and back", this saves
145 the ceremony::
147 point_transformer = SimpleTransformer(
148 "point", Point,
149 to_dict=lambda p: {"x": p.x, "y": p.y},
150 from_dict=lambda d: Point(d["x"], d["y"]),
151 )
153 The callables receive and return plain values; nested values are *not*
154 transformed automatically. Subclass :class:`CustomTransformer` directly when
155 the encoding needs to recurse through the transformer, or to write bytes
156 out-of-line.
157 """
159 def __init__(
160 self,
161 name: str,
162 type_: type,
163 *,
164 to_dict: Callable[[Any], dict[str, Any]],
165 from_dict: Callable[[dict[str, Any]], Any],
166 subtypes: bool = False,
167 ) -> None:
168 """Build a transformer for *type_* from *to_dict* and *from_dict*.
170 :param name: Unique discriminator written as the ``"type"`` field.
171 :param type_: The type this transformer handles.
172 :param to_dict: Callable taking an instance and returning a plain dict.
173 :param from_dict: Callable taking that dict and returning an instance.
174 :param subtypes: When true, also handle subclasses of *type_*.
175 """
176 self._name = name
177 self._type = type_
178 self._to_dict = to_dict
179 self._from_dict = from_dict
180 self._subtypes = subtypes
182 @property
183 def name(self) -> str:
184 """Return the transformer's discriminator name."""
185 return self._name
187 def to_dict(self, transformer: "Transformer", o: object) -> dict[str, Any]:
188 """Encode *o* by delegating to the supplied callable."""
189 return self._to_dict(o)
191 def from_dict(self, transformer: "Transformer", d: dict[str, Any]) -> object:
192 """Decode *d* by delegating to the supplied callable."""
193 return self._from_dict(d)
195 @property
196 def supported_direct_types(self) -> Iterable[type]:
197 """Return the handled type, unless subtype matching was requested."""
198 return [] if self._subtypes else [self._type]
200 @property
201 def supported_subtypes(self) -> Iterable[Any]:
202 """Return the handled base type when subtype matching was requested."""
203 return [self._type] if self._subtypes else []
206class Transformable(ABC):
207 """Abstract base class for objects that can transform themselves."""
209 @abstractmethod
210 def to_dict(self, transformer: "Transformer") -> dict[str, Any]:
211 """Convert this object to dictionary representation."""
212 # pragma: no cover
214 @classmethod
215 @abstractmethod
216 def from_dict(cls, transformer: "Transformer", d: dict[str, Any]) -> object:
217 """Reconstruct object from dictionary representation."""
218 # pragma: no cover
221class Transformer:
222 """Main transformer class for object serialization and deserialization."""
224 def __init__(self, *, strict: bool = True) -> None:
225 """Initialize transformer with strict mode setting."""
226 self.strict = strict
228 self._direct_type_map: dict[type, CustomTransformer] = {}
229 self._subtype_order: list[type] = []
230 self._subtype_map: dict[type, CustomTransformer] = {}
231 self._transformers: dict[str, CustomTransformer] = {}
232 self._transformable_types: dict[str, type[Transformable]] = {}
233 self._attrs_types: dict[str, type] = {}
234 self._dataclass_types: dict[str, type] = {}
236 def register(self, t: CustomTransformer | type[Transformable] | type) -> None:
237 """Register a transformer, transformable type, or regular type."""
238 if isinstance(t, CustomTransformer):
239 self.register_transformer(t)
240 elif isinstance(t, type) and issubclass(t, Transformable):
241 self.register_transformable(t)
242 elif HAS_ATTRS and isinstance(t, type) and attrs.has(t):
243 self.register_attrs(t)
244 elif isinstance(t, type) and dataclasses.is_dataclass(t):
245 self.register_dataclass(t)
246 else:
247 msg = f"Unable to register {t}"
248 raise ValueError(msg)
250 def register_transformer(self, transformer: CustomTransformer) -> None:
251 """Register a custom transformer for specific types."""
252 if transformer.name in self._transformers:
253 msg = f"A transformer named {transformer.name!r} is already registered"
254 raise DuplicateRegistrationError(msg)
255 for type_ in transformer.supported_direct_types:
256 if type_ in self._direct_type_map:
257 msg = f"Type {type_!r} is already handled directly by transformer {self._direct_type_map[type_].name!r}"
258 raise DuplicateRegistrationError(msg)
259 for type_ in transformer.supported_subtypes:
260 if type_ in self._subtype_map:
261 msg = f"Subtype {type_!r} is already handled by transformer {self._subtype_map[type_].name!r}"
262 raise DuplicateRegistrationError(msg)
264 self._transformers[transformer.name] = transformer
266 for type_ in transformer.supported_direct_types:
267 self._direct_type_map[type_] = transformer
269 contains_supported_subtypes = False
270 for type_ in transformer.supported_subtypes:
271 contains_supported_subtypes = True
272 self._subtype_map[type_] = transformer
273 if contains_supported_subtypes:
274 self._subtype_order = order_classes(self._subtype_map.keys())
276 def register_transformable(self, transformable_type: type[Transformable]) -> None:
277 """Register a transformable type that can serialize itself."""
278 name = transformable_type.__name__
279 if name in self._transformable_types:
280 msg = f"A Transformable class named {name!r} is already registered"
281 raise DuplicateRegistrationError(msg)
282 self._transformable_types[name] = transformable_type
284 def register_attrs(self, attrs_type: type) -> None:
285 """Register an attrs-decorated class for serialization."""
286 name = attrs_type.__name__
287 if name in self._attrs_types:
288 msg = f"An attrs class named {name!r} is already registered"
289 raise DuplicateRegistrationError(msg)
290 self._attrs_types[name] = attrs_type
292 def register_dataclass(self, dataclass_type: type) -> None:
293 """Register a dataclass for serialization."""
294 name = dataclass_type.__name__
295 if name in self._dataclass_types:
296 msg = f"A dataclass named {name!r} is already registered"
297 raise DuplicateRegistrationError(msg)
298 self._dataclass_types[name] = dataclass_type
300 def get_transformer_for_obj(self, obj: object) -> CustomTransformer | None:
301 """Get the appropriate transformer for a given object."""
302 transformer = self._direct_type_map.get(type(obj))
303 if transformer is not None:
304 return transformer
305 for tp in self._subtype_order:
306 if isinstance(obj, tp):
307 return self._subtype_map[tp]
308 return None
310 def get_transformer_for_name(self, name: str) -> CustomTransformer | None:
311 """Get a transformer by its registered name."""
312 transformer = self._transformers.get(name)
313 return transformer
315 # ------------------------------------------------------------------
316 # Out-of-line bytes.
317 #
318 # A transformer is *offered* a blob sink rather than being asked to
319 # implement a second encoding: it calls offer_blob() to ask whether this
320 # save wants bytes out of line, and put_blob() to hand them over. Nothing
321 # about to_dict's signature changes, so a transformer written before any of
322 # this existed keeps working --- it never calls offer_blob, so it always
323 # inlines, which is exactly what it did before.
324 #
325 # The sink lives in a ContextVar rather than being threaded through every
326 # call. That keeps user code free of plumbing, and makes concurrent use safe
327 # for free, which matters because loman runs nodes on executors.
328 # ------------------------------------------------------------------
330 @contextlib.contextmanager
331 def writing(
332 self,
333 sink: Any,
334 profile: Any = None,
335 node: str | None = None,
336 tags: frozenset[str] = frozenset(),
337 store: str | None = None,
338 ) -> Iterator[None]:
339 """Scope in which ``offer_blob`` and ``put_blob`` write to *sink*.
341 :param node: Node key, for the blob table and for glob selectors.
342 :param tags: The node's tags, for ``tag:`` selectors.
343 :param store: The store named on the node itself, which a profile
344 override for this node can replace.
345 """
346 token = _WRITE_SCOPE.set(_WriteScope(sink=sink, profile=profile, node=node, tags=tags, store=store))
347 try:
348 yield
349 finally:
350 _WRITE_SCOPE.reset(token)
352 @contextlib.contextmanager
353 def reading(self, source: Any) -> Iterator[None]:
354 """Scope in which ``get_blob`` resolves references against *source*."""
355 token = _READ_SCOPE.set(source)
356 try:
357 yield
358 finally:
359 _READ_SCOPE.reset(token)
361 def offer_blob(self, nbytes: int | None = None) -> bool:
362 """Return whether a value of *nbytes* should be written out of line.
364 Outside a write scope, or in a container with no blob storage, this is
365 always false --- so a transformer never has to know which container it
366 is writing into.
368 :param nbytes: Estimated encoded size. ``None`` means "large, but I
369 cannot say how large", and is taken at its word.
370 """
371 scope = _WRITE_SCOPE.get()
372 if scope is None or scope.profile is None:
373 return False
374 # Asks the writer about *this node's* store, not just whether any store
375 # exists. A node routed to an external store can go out of line even in
376 # the single-document container, which is how a readable manifest can
377 # sit alongside data held in S3.
378 store = self.blob_store_name()
379 if not scope.sink.can_store(store):
380 if store is not None:
381 # The node asked for a specific store and it was not supplied.
382 # Falling back to inline would put the data in the file while
383 # the caller believed it had gone to their bucket, so this is an
384 # error rather than a quiet substitution.
385 from loman.exception import SerializationError
387 msg = (
388 f"Node {scope.node!r} is routed to blob store {store!r}, which was not supplied. "
389 f"Pass it as save(..., stores={{{store!r}: ...}}), or override the routing for "
390 "this save with a profile."
391 )
392 raise SerializationError(msg)
393 return False
394 return bool(scope.profile.wants_blob(nbytes))
396 def blob_store_name(self) -> str | None:
397 """Return the store the value being encoded should be written to.
399 The node's own declaration is the default; a profile override matching
400 this node replaces it, so the same computation can be saved to S3 in
401 production and to a plain container in a test.
402 """
403 scope = _WRITE_SCOPE.get()
404 if scope is None:
405 return None
406 return self.blob_setting("store", scope.store)
408 def put_blob(
409 self,
410 payload: Any,
411 *,
412 codec: str,
413 compressible: bool = True,
414 dedupe_on: Any = None,
415 ) -> dict[str, int]:
416 """Store *payload* out of line and return the reference to embed.
418 :param payload: ``bytes``, a buffer, or a callable taking a binary file
419 object. The callable form lets an encoder stream straight into the
420 container instead of building the whole payload in memory first.
421 :param codec: How the bytes are encoded, e.g. ``"npy"``. Recorded in the
422 blob table as the file extension and as metadata for other tools; it
423 is not what dispatch keys off on the way back in.
424 :param compressible: Pass false for a payload that already compresses
425 itself, such as parquet. Skips the compression probe entirely, which
426 is how double-compression is prevented.
427 :param dedupe_on: The object being stored, when two nodes holding the
428 same object should share one blob. Pass the object itself, not its
429 ``id()``: the store keeps a reference to it, which is what stops a
430 short-lived temporary's id from being reused by the next one and
431 silently deduplicating two unrelated values onto one blob.
432 """
433 scope = _WRITE_SCOPE.get()
434 if scope is None:
435 msg = "put_blob() called outside a write scope; use Transformer.writing()"
436 raise RuntimeError(msg)
437 from .blobs import blob_ref
439 blob_id = scope.sink.put(
440 payload,
441 codec=codec,
442 node=scope.node,
443 store=self.blob_store_name(),
444 compressible=compressible,
445 dedupe_on=dedupe_on,
446 )
447 return blob_ref(blob_id)
449 def blob_setting(self, name: str, default: Any = None) -> Any:
450 """Return a profile setting for the value currently being encoded.
452 Per-node overrides win over the profile's own value, so one save can
453 treat some nodes differently from the rest. Outside a write scope the
454 default is returned, which is what keeps transformers working when they
455 are driven directly.
456 """
457 scope = _WRITE_SCOPE.get()
458 if scope is None or scope.profile is None:
459 return default
460 overrides = scope.profile.settings_for(scope.node, scope.tags)
461 if name in overrides:
462 return overrides[name]
463 return getattr(scope.profile, name, default)
465 def get_blob(self, ref: dict[str, int]) -> bytes:
466 """Return the bytes for the blob reference *ref*."""
467 source = _READ_SCOPE.get()
468 if source is None:
469 msg = "This value's data is stored out of line, but no blob store is open for reading"
470 raise RuntimeError(msg)
471 from .blobs import BLOB_REF_KEY
473 return source.get(ref[BLOB_REF_KEY])
475 def to_dict(self, o: object) -> Any:
476 """Convert an object to a serializable dictionary representation.
478 Split in two along the line that matters. This half handles values whose
479 type is *exactly* a builtin, which is the overwhelming majority of what a
480 graph holds and so the path worth keeping short; everything else goes to
481 :meth:`_to_dict_object`.
483 The checks are exact-type, not ``isinstance``, deliberately: numpy
484 scalars and ``IntEnum`` members are instances of ``int`` or ``float``, so
485 an ``isinstance`` fast path would return them as bare numbers and
486 silently discard the type. Falling through is what lets the registered
487 transformers see them.
488 """
489 type_ = type(o)
490 if type_ is str or o is None or o is True or o is False or type_ is int:
491 return o
492 if type_ is float:
493 # cast: an exact-type check carries no narrowing for a type checker.
494 number = cast("float", o)
495 return number if math.isfinite(number) else {KEY_TYPE: TYPENAME_FLOAT, KEY_VALUES: _nonfinite_name(number)}
496 if type_ is tuple:
497 return {KEY_TYPE: TYPENAME_TUPLE, KEY_VALUES: [self.to_dict(x) for x in cast("tuple[Any, ...]", o)]}
498 if type_ is list:
499 return [self.to_dict(x) for x in cast("list[Any]", o)]
500 if type_ is dict:
501 return self._dict_to_dict(cast("dict[Any, Any]", o))
502 return self._to_dict_object(o)
504 def _to_dict_object(self, o: object) -> Any:
505 """Encode a value that is not exactly a builtin scalar or container.
507 Order is significant. Registered transformers come first so an
508 explicitly registered type --- ``NodeKey``, say --- wins over the generic
509 attrs and dataclass handling that would otherwise also claim it.
510 """
511 if self.get_transformer_for_obj(o) is not None:
512 return self._to_dict_transformer(o)
513 if isinstance(o, Transformable):
514 return {KEY_TYPE: TYPENAME_TRANSFORMABLE, KEY_CLASS: type(o).__name__, KEY_DATA: o.to_dict(self)}
515 if HAS_ATTRS and attrs.has(type(o)):
516 return self._attrs_to_dict(o)
517 if dataclasses.is_dataclass(o) and not isinstance(o, type):
518 return self._dataclass_to_dict(o)
520 # Subclasses of the builtin containers, reached only when nothing above
521 # claimed them: a namedtuple, an OrderedDict, a list subclass. Encoded as
522 # the plain builtin form, which is what they did before exact-type
523 # dispatch was introduced.
524 if isinstance(o, tuple):
525 return {KEY_TYPE: TYPENAME_TUPLE, KEY_VALUES: [self.to_dict(x) for x in o]}
526 if isinstance(o, list):
527 return [self.to_dict(x) for x in o]
528 if isinstance(o, dict):
529 return self._dict_to_dict(o)
530 return self._to_dict_transformer(o)
532 def _dict_to_dict(self, o: dict[Any, Any]) -> dict[str, Any]:
533 """Convert a dictionary to serializable form.
535 A dict whose keys are all plain strings, none of them reserved, is
536 written as a JSON object --- the readable form, and the overwhelmingly
537 common case. Anything else is written as a list of encoded key/value
538 pairs, because a JSON object key can only be a string: the previous
539 encoding turned ``{1: 'a'}`` into ``{'1': 'a'}`` and handed back a
540 different dict than it was given.
541 """
542 if all(type(k) is str and k not in RESERVED_KEYS for k in o):
543 return {k: self.to_dict(v) for k, v in o.items()}
544 return {
545 KEY_TYPE: TYPENAME_DICT,
546 KEY_ITEMS: [[self.to_dict(k), self.to_dict(v)] for k, v in o.items()],
547 }
549 def _attrs_to_dict(self, o: object) -> dict[str, Any]:
550 """Convert an attrs object to serializable dictionary form."""
551 data: dict[str, Any] = {}
552 for a in o.__attrs_attrs__: # type: ignore[attr-defined]
553 data[a.name] = self.to_dict(o.__getattribute__(a.name))
554 res: dict[str, Any] = {KEY_TYPE: TYPENAME_ATTRS, KEY_CLASS: type(o).__name__}
555 if len(data) > 0:
556 res[KEY_DATA] = data
557 return res
559 def _dataclass_to_dict(self, o: object) -> dict[str, Any]:
560 """Convert a dataclass object to serializable dictionary form."""
561 data: dict[str, Any] = {}
562 for f in dataclasses.fields(o): # type: ignore[arg-type]
563 data[f.name] = self.to_dict(getattr(o, f.name))
564 res: dict[str, Any] = {KEY_TYPE: TYPENAME_DATACLASS, KEY_CLASS: type(o).__name__}
565 if len(data) > 0:
566 res[KEY_DATA] = data
567 return res
569 def _to_dict_transformer(self, o: object) -> dict[str, Any] | None:
570 """Convert an object using a registered custom transformer."""
571 transformer = self.get_transformer_for_obj(o)
572 if transformer is None:
573 if self.strict:
574 msg = f"Could not transform object of type {type(o).__name__}"
575 raise UntransformableTypeError(msg)
576 else:
577 return None
578 d = transformer.to_dict(self, o)
579 d[KEY_TYPE] = transformer.name
580 return d
582 def from_dict(self, d: Any) -> Any:
583 """Convert a dictionary representation back to the original object."""
584 if isinstance(d, str) or d is None or d is True or d is False or isinstance(d, (int, float)):
585 return d
586 elif isinstance(d, list):
587 return [self.from_dict(x) for x in d]
588 elif isinstance(d, dict):
589 type_ = d.get(KEY_TYPE)
590 if type_ is None:
591 return {k: self.from_dict(v) for k, v in d.items()}
592 elif type_ == TYPENAME_FLOAT:
593 return _NAME_TO_FLOAT[d[KEY_VALUES]]
594 elif type_ == TYPENAME_TUPLE:
595 return tuple(self.from_dict(x) for x in d[KEY_VALUES])
596 elif type_ == TYPENAME_DICT:
597 return self._from_escaped_dict(d)
598 elif type_ == TYPENAME_TRANSFORMABLE:
599 return self._from_dict_transformable(d)
600 elif type_ == TYPENAME_ATTRS:
601 return self._from_attrs(d)
602 elif type_ == TYPENAME_DATACLASS:
603 return self._from_dataclass(d)
604 else:
605 return self._from_dict_transformer(type_, d)
606 else:
607 msg = "Unable to determine object type from dictionary"
608 raise ValueError(msg)
610 def _from_escaped_dict(self, d: dict[str, Any]) -> dict[Any, Any]:
611 """Reconstruct a dict written in the escaped form.
613 Two layouts are accepted. ``items`` is the current one, a list of encoded
614 key/value pairs that supports keys of any transformable type. ``data`` is
615 the older layout, a JSON object that could only ever hold string keys;
616 files written before this change still use it.
617 """
618 if KEY_ITEMS in d:
619 return {self.from_dict(k): self.from_dict(v) for k, v in d[KEY_ITEMS]}
620 return {k: self.from_dict(v) for k, v in d[KEY_DATA].items()}
622 def _from_dict_transformable(self, d: dict[str, Any]) -> object:
623 """Reconstruct a Transformable object from dictionary form."""
624 classname = d[KEY_CLASS]
625 cls = self._transformable_types.get(classname)
626 if cls is None:
627 if self.strict:
628 msg = f"Unable to transform Transformable object of class {classname}"
629 raise UnrecognizedTypeError(msg)
630 else:
631 return MissingObject()
632 else:
633 return cls.from_dict(self, d[KEY_DATA])
635 def _from_attrs(self, d: dict[str, Any]) -> object:
636 """Reconstruct an attrs object from dictionary form."""
637 if not HAS_ATTRS: # pragma: no cover
638 if self.strict:
639 msg = "attrs package not installed"
640 raise UnrecognizedTypeError(msg)
641 return MissingObject()
642 cls = self._attrs_types.get(d[KEY_CLASS])
643 if cls is None:
644 if self.strict:
645 msg = f"Unable to create attrs object of type {cls}"
646 raise UnrecognizedTypeError(msg)
647 else:
648 return MissingObject()
649 else:
650 kwargs: dict[str, Any] = {}
651 if KEY_DATA in d:
652 for key, value in d[KEY_DATA].items():
653 kwargs[key] = self.from_dict(value)
654 return cls(**kwargs)
656 def _from_dataclass(self, d: dict[str, Any]) -> object:
657 """Reconstruct a dataclass object from dictionary form."""
658 cls = self._dataclass_types.get(d[KEY_CLASS])
659 if cls is None:
660 if self.strict:
661 msg = f"Unable to create dataclass object of type {cls}"
662 raise UnrecognizedTypeError(msg)
663 else:
664 return MissingObject()
665 else:
666 kwargs: dict[str, Any] = {}
667 if KEY_DATA in d:
668 for key, value in d[KEY_DATA].items():
669 kwargs[key] = self.from_dict(value)
670 return cls(**kwargs)
672 def _from_dict_transformer(self, type_: str, d: dict[str, Any]) -> object:
673 """Reconstruct an object using a registered custom transformer."""
674 transformer = self.get_transformer_for_name(type_)
675 if transformer is None:
676 if self.strict:
677 msg = f"Unable to transform object of type {type_}"
678 raise UnrecognizedTypeError(msg)
679 else:
680 return MissingObject()
681 return transformer.from_dict(self, d)
684class NdArrayTransformer(CustomTransformer):
685 """Transformer for NumPy ndarray objects."""
687 @property
688 def name(self) -> str:
689 """Return transformer name."""
690 return "ndarray"
692 def to_dict(self, transformer: "Transformer", o: object) -> dict[str, Any]:
693 """Convert numpy array to dictionary with shape, dtype, and data.
695 The shape and dtype stay inline whichever way the data goes, so the
696 manifest still describes every array in the graph without a single blob
697 being decoded.
698 """
699 assert isinstance(o, np.ndarray) # noqa: S101
700 head: dict[str, Any] = {"shape": list(o.shape), "dtype": o.dtype.str}
701 if o.dtype.hasobject:
702 # An object array's elements are arbitrary Python values; .npy would
703 # have to pickle them, so it stays on the inline path.
704 head["data"] = transformer.to_dict(o.ravel().tolist()) # type: ignore[arg-type]
705 return head
706 if transformer.offer_blob(nbytes=o.nbytes):
707 head["encoding"] = "npy"
708 head["data"] = transformer.put_blob(
709 lambda f: np.save(f, o, allow_pickle=False),
710 codec="npy",
711 dedupe_on=o,
712 )
713 else:
714 head["data"] = transformer.to_dict(o.ravel().tolist()) # type: ignore[arg-type]
715 return head
717 def from_dict(self, transformer: "Transformer", d: dict[str, Any]) -> object:
718 """Reconstruct numpy array from dictionary."""
719 if d.get("encoding") == "npy":
720 import io
722 return np.load(io.BytesIO(transformer.get_blob(d["data"])), allow_pickle=False)
723 return np.array(transformer.from_dict(d["data"]), d["dtype"]).reshape(d["shape"])
725 @property
726 def supported_direct_types(self) -> Iterable[type]:
727 """Return supported numpy array types."""
728 return [np.ndarray]
731class EnumTransformer(CustomTransformer):
732 """Transformer for Enum subclasses.
734 Enum classes must be registered via :meth:`register_enum` before use.
735 """
737 def __init__(self) -> None:
738 """Initialise with an empty enum registry."""
739 self._registry: dict[str, type[Enum]] = {}
741 def register_enum(self, enum_class: type[Enum]) -> None:
742 """Register an enum class so its members can be deserialized."""
743 self._registry[enum_class.__qualname__] = enum_class
745 @property
746 def name(self) -> str:
747 """Return transformer name."""
748 return "enum"
750 def to_dict(self, transformer: "Transformer", o: object) -> dict[str, Any]:
751 """Convert an Enum member to a dict with class qualname and member name."""
752 assert isinstance(o, Enum) # noqa: S101
753 return {"enum_class": type(o).__qualname__, "value": o.name}
755 def from_dict(self, transformer: "Transformer", d: dict[str, Any]) -> object:
756 """Reconstruct an Enum member from its serialized form."""
757 enum_class = self._registry.get(d["enum_class"])
758 if enum_class is None:
759 msg = f"Unknown enum class: {d['enum_class']!r}. Register it with EnumTransformer.register_enum()."
760 raise UnrecognizedTypeError(msg)
761 return enum_class[d["value"]]
763 @property
764 def supported_subtypes(self) -> Iterable[type]:
765 """Handle all Enum subclasses."""
766 return [Enum]
769class FunctionRefTransformer(CustomTransformer):
770 """Transformer for importable callables (module-level functions and methods).
772 Lambdas and closures (whose ``__qualname__`` contains ``<lambda>`` or
773 ``<locals>``) are explicitly rejected with a :class:`ValueError`.
774 """
776 @property
777 def name(self) -> str:
778 """Return transformer name."""
779 return "func_ref"
781 def to_dict(self, transformer: "Transformer", o: object) -> dict[str, Any]:
782 """Serialize a callable as its module path and qualname."""
783 if not callable(o):
784 msg = f"Object {o!r} is not callable"
785 raise TypeError(msg)
786 qualname = getattr(o, "__qualname__", None)
787 module = getattr(o, "__module__", None)
788 if qualname is None or module is None:
789 msg = f"Cannot serialize {o!r}: missing __qualname__ or __module__"
790 raise ValueError(msg)
791 if "<lambda>" in qualname:
792 msg = f"Cannot serialize lambda function {o!r}: lambdas are not importable"
793 raise ValueError(msg)
794 if "<locals>" in qualname:
795 msg = f"Cannot serialize closure/local function {o!r}: non-importable"
796 raise ValueError(msg)
797 # Verify the callable is actually reachable via import before committing.
798 try:
799 mod = importlib.import_module(module)
800 obj: Any = mod
801 for part in qualname.split("."):
802 obj = getattr(obj, part)
803 if obj is not o:
804 msg = f"Cannot serialize {o!r}: import round-trip returned a different object"
805 raise ValueError(msg)
806 except (ImportError, AttributeError) as exc:
807 msg = f"Cannot serialize {o!r}: not importable ({exc})"
808 raise ValueError(msg) from exc
809 return {"module": module, "qualname": qualname}
811 def from_dict(self, transformer: "Transformer", d: dict[str, Any]) -> object:
812 """Reconstruct a callable from its module path and qualname."""
813 module = importlib.import_module(d["module"])
814 obj: Any = module
815 for part in d["qualname"].split("."):
816 obj = getattr(obj, part)
817 return obj
819 @property
820 def supported_direct_types(self) -> Iterable[type]:
821 """Register the built-in function types explicitly handled."""
822 # We use supported_subtypes for the broad callable match instead,
823 # but we must list at least one concrete type here to help dispatch.
824 # The broad subtype match on Callable covers everything callable.
825 return []
827 @property
828 def supported_subtypes(self) -> Iterable[Any]:
829 """Match all callables via Callable ABC."""
830 return [Callable]
833class DefinitionMethodTransformer(CustomTransformer):
834 """Transformer for calc nodes declared on a ``@ComputationFactory`` class.
836 A computation built from a class stores each calc node's function as a method
837 *bound to the definition object* --- the instance the factory created. That
838 has no importable path of its own:
839 :class:`FunctionRefTransformer` looks up ``module.Portfolio.signal``, but
840 after decoration the name ``Portfolio`` refers to the factory function rather
841 than to the class, so the lookup fails.
843 The consequence, before this transformer existed, was that the library's own
844 primary idiom produced graphs whose functions were silently dropped. Values
845 reloaded; nodes could never recompute.
847 The route back is through the factory. ``functools.wraps(cls)`` leaves the
848 class on the factory as ``__wrapped__``, so the method is stored as the class
849 name plus the method name, and rebuilt by importing the module, recovering
850 the class, and binding the method to a definition object.
852 .. note::
853 That definition object is a **new** instance, not the one the graph was
854 built with. For the ordinary case --- a class whose body only declares
855 nodes --- the two are indistinguishable. A class whose ``__init__``
856 computes state, or whose methods mutate ``self`` at run time, will not
857 see that state after a round-trip. A class that cannot be constructed
858 without arguments cannot be restored at all, and its node falls back to
859 being stored as ``null`` with a warning.
861 One instance is created per class per transformer, so all of a
862 computation's nodes share a definition object, as they did originally.
863 """
865 def __init__(self) -> None:
866 """Start with no definition objects built."""
867 self._instances: dict[type, object] = {}
869 @property
870 def name(self) -> str:
871 """Return transformer name."""
872 return "definition_method"
874 def to_dict(self, transformer: "Transformer", o: object) -> dict[str, Any]:
875 """Serialize a bound method as its class name and method name."""
876 func = getattr(o, "__func__", None)
877 if func is None:
878 msg = f"Cannot serialize {o!r}: not a bound method"
879 raise ValueError(msg)
881 qualname = getattr(func, "__qualname__", "") or ""
882 module = getattr(func, "__module__", None)
883 if module is None or "." not in qualname:
884 msg = f"Cannot serialize {o!r}: no class-qualified name to resolve it by"
885 raise ValueError(msg)
886 if "<locals>" in qualname:
887 msg = f"Cannot serialize {o!r}: defined inside a function, so it is not importable"
888 raise ValueError(msg)
890 class_qualname, _, method = qualname.rpartition(".")
892 # Resolve now rather than on load. A method that cannot be rebuilt should
893 # be reported while the graph is still in front of the person saving it,
894 # not when someone else opens the file.
895 rebuilt = self._resolve(module, class_qualname, method)
896 if getattr(rebuilt, "__func__", None) is not func:
897 msg = f"Cannot serialize {o!r}: {module}.{qualname} does not resolve back to this method"
898 raise ValueError(msg)
900 return {"module": module, "class_qualname": class_qualname, "method": method}
902 def from_dict(self, transformer: "Transformer", d: dict[str, Any]) -> object:
903 """Rebuild the bound method from its class and method names."""
904 return self._resolve(d["module"], d["class_qualname"], d["method"])
906 def _resolve(self, module: str, class_qualname: str, method: str) -> object:
907 """Return the bound method for *method* on the definition class."""
908 cls = self._definition_class(module, class_qualname)
909 instance = self._instances.get(cls)
910 if instance is None:
911 try:
912 instance = cls()
913 except TypeError as exc:
914 msg = (
915 f"Cannot rebuild {class_qualname}: its definition class cannot be constructed "
916 f"without arguments ({exc})"
917 )
918 raise ValueError(msg) from exc
919 self._instances[cls] = instance
921 bound = getattr(instance, method, None)
922 if bound is None or not callable(bound):
923 msg = f"Cannot rebuild {class_qualname}.{method}: no such method on the definition class"
924 raise ValueError(msg)
925 return bound
927 @staticmethod
928 def _definition_class(module: str, class_qualname: str) -> type:
929 """Return the class named by *class_qualname*, seeing through the factory."""
930 try:
931 obj: Any = importlib.import_module(module)
932 for part in class_qualname.split("."):
933 obj = getattr(obj, part)
934 except (ImportError, AttributeError) as exc:
935 msg = f"Cannot resolve {module}.{class_qualname}: {exc}"
936 raise ValueError(msg) from exc
938 # After @ComputationFactory the name refers to the factory function;
939 # functools.wraps left the class on it as __wrapped__.
940 if not isinstance(obj, type):
941 obj = getattr(obj, "__wrapped__", obj)
942 if not isinstance(obj, type):
943 # ValueError, not TypeError: the complaint is about what this name
944 # resolved to, not about an argument the caller passed. Every other
945 # unresolvable-callable path raises ValueError too, and the callers
946 # that turn this into a warning key off that.
947 msg = f"Cannot resolve {module}.{class_qualname}: it is not a class or a computation factory"
948 raise ValueError(msg) # noqa: TRY004
949 return obj
951 @property
952 def supported_subtypes(self) -> Iterable[Any]:
953 """Match bound methods, which are more specific than plain callables."""
954 return [types.MethodType]
957class DillFunctionTransformer(CustomTransformer):
958 """Transformer that serializes any callable — including lambdas and closures — using dill.
960 The callable is serialized with :func:`dill.dumps` and the resulting bytes
961 are stored as a base64-encoded string inside the JSON document. On load the
962 bytes are decoded and passed to :func:`dill.loads`.
964 .. note::
965 The embedded dill blob is **not** portable across Python versions and
966 shares the same stability caveats as :meth:`~loman.Computation.write_dill`.
967 Register this transformer when convenient lambda/closure round-trips matter
968 more than portability.
970 Example::
972 from loman import Computation, ComputationSerializer
973 from loman.serialization import DillFunctionTransformer
975 s = ComputationSerializer(use_dill_for_functions=True)
976 comp = Computation()
977 comp.add_node('a', value=1)
978 comp.add_node('b', lambda a: a + 1)
979 comp.compute_all()
980 comp.write_json('comp.json', serializer=s)
981 comp2 = Computation.read_json('comp.json', serializer=s)
982 assert comp2.v.b == 2
983 """
985 @property
986 def name(self) -> str:
987 """Return transformer name."""
988 return "dill_func"
990 def to_dict(self, transformer: "Transformer", o: object) -> dict[str, Any]:
991 """Serialize a callable to a base64-encoded dill blob."""
992 import base64
994 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
996 if not callable(o):
997 msg = f"Object {o!r} is not callable"
998 raise TypeError(msg)
999 blob = dill.dumps(o)
1000 return {"blob": base64.b64encode(blob).decode("ascii")}
1002 def from_dict(self, transformer: "Transformer", d: dict[str, Any]) -> object:
1003 """Reconstruct a callable from a base64-encoded dill blob."""
1004 import base64
1006 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
1008 blob = base64.b64decode(d["blob"].encode("ascii"))
1009 return dill.loads(blob) # noqa: S301 # nosec B301
1011 @property
1012 def supported_direct_types(self) -> Iterable[type]:
1013 """No direct type matches — rely on subtype matching."""
1014 return []
1016 @property
1017 def supported_subtypes(self) -> Iterable[Any]:
1018 """Match all callables via Callable ABC."""
1019 return [Callable]
1022def _encode_frame_as_parquet(transformer: "Transformer", frame: "pd.DataFrame") -> dict[str, Any] | None:
1023 """Encode *frame* as a parquet blob, or return ``None`` to use the default path.
1025 Parquet is opt-in via ``frame_encoding="parquet"`` on the profile. It buys
1026 columnar compression and a file other tools can read, at the cost of an
1027 optional pyarrow dependency.
1029 Returns ``None`` --- meaning "not this way" --- when parquet was not asked
1030 for, when the value is too small to be worth a blob, or when pyarrow cannot
1031 represent this particular frame. That last case is why the conversion is
1032 attempted before anything is written: a frame with duplicate column names or
1033 an exotic dtype should fall back to an encoding that works, not fail the
1034 save.
1035 """
1036 if transformer.blob_setting("frame_encoding", "npy") != "parquet":
1037 return None
1039 estimated = int(frame.memory_usage(deep=False).sum())
1040 if not transformer.offer_blob(nbytes=estimated):
1041 return None
1043 try:
1044 from loman._extras import require
1046 pa = require("pyarrow", "efficient")
1047 pq = require("pyarrow.parquet", "efficient")
1048 table = pa.Table.from_pandas(frame, preserve_index=True)
1049 except Exception:
1050 return None
1052 def write(f: Any) -> None:
1053 pq.write_table(table, f, compression="zstd")
1055 return {
1056 "shape": list(frame.shape),
1057 "encoding": "parquet",
1058 # Parquet already compresses; compressing the blob again would cost time
1059 # and achieve nothing.
1060 "data": transformer.put_blob(write, codec="parquet", compressible=False, dedupe_on=frame),
1061 }
1064def _encode_column(transformer: "Transformer", column: "pd.Series") -> Any:
1065 """Encode one DataFrame column, as an array where that is lossless.
1067 A column backed by a plain numpy dtype is handed to the ndarray transformer,
1068 which means it inherits out-of-line storage for free --- this is what keeps a
1069 large numeric frame from being written one decimal string at a time.
1071 Extension dtypes (categorical, nullable integers, timezone-aware datetimes,
1072 pandas strings) are *not* sent that way. Their numpy representation loses the
1073 thing that makes them distinct --- a tz-aware column flattens to naive UTC,
1074 and reading it back as a local wall time would silently shift every value ---
1075 so they go through the element-wise path and are restored from their
1076 recorded dtype.
1077 """
1078 dtype = column.dtype
1079 if isinstance(dtype, np.dtype) and not dtype.hasobject:
1080 return transformer.to_dict(column.to_numpy())
1081 return transformer.to_dict(column.tolist())
1084class DataFrameTransformer(CustomTransformer):
1085 """Transformer for :class:`pandas.DataFrame` objects."""
1087 @property
1088 def name(self) -> str:
1089 """Return transformer name."""
1090 return "dataframe"
1092 def to_dict(self, transformer: "Transformer", o: object) -> dict[str, Any]:
1093 """Serialize a DataFrame using split orientation.
1095 Columns are encoded per column rather than through ``o.values``, which
1096 would push every cell through one shared dtype --- object for a frame
1097 mixing numbers and strings --- and lose the per-column types on the way.
1098 Both axes go through the transformer as indexes, so a ``MultiIndex`` or a
1099 ``DatetimeIndex`` on either axis survives the round-trip.
1100 """
1101 assert isinstance(o, pd.DataFrame) # noqa: S101
1102 parquet = _encode_frame_as_parquet(transformer, o)
1103 if parquet is not None:
1104 return parquet
1105 return {
1106 "columns": transformer.to_dict(o.columns),
1107 "index": transformer.to_dict(o.index),
1108 "data": [_encode_column(transformer, o.iloc[:, i]) for i in range(o.shape[1])],
1109 "orient": "columns",
1110 "dtypes": [str(dtype) for dtype in o.dtypes],
1111 }
1113 def from_dict(self, transformer: "Transformer", d: dict[str, Any]) -> object:
1114 """Reconstruct a DataFrame from its serialized form."""
1115 # A parquet blob carries its own schema, columns and index in the file
1116 # footer, so it is decoded before anything else is read from the entry.
1117 if d.get("encoding") == "parquet":
1118 import io
1120 from loman._extras import require
1122 pq = require("pyarrow.parquet", "efficient")
1123 return pq.read_table(io.BytesIO(transformer.get_blob(d["data"]))).to_pandas()
1125 columns = transformer.from_dict(d["columns"])
1126 index = transformer.from_dict(d["index"])
1127 dtypes = d.get("dtypes", {})
1129 if d.get("orient") == "columns":
1130 data = {i: transformer.from_dict(col) for i, col in enumerate(d["data"])}
1131 df = pd.DataFrame(data, index=index)
1132 df.columns = pd.Index(columns) if not isinstance(columns, pd.Index) else columns
1133 for i, dtype in enumerate(dtypes):
1134 with contextlib.suppress(ValueError, TypeError):
1135 df.isetitem(i, df.iloc[:, i].astype(dtype))
1136 return df
1138 # Row-major layout written before the per-column encoding existed.
1139 df = pd.DataFrame(transformer.from_dict(d["data"]), columns=columns, index=index)
1140 for col, dtype in (dtypes or {}).items():
1141 with contextlib.suppress(ValueError, TypeError): # pragma: no cover
1142 df[col] = df[col].astype(dtype)
1143 return df
1145 @property
1146 def supported_direct_types(self) -> Iterable[type]:
1147 """Return supported pandas DataFrame type."""
1148 return [pd.DataFrame]
1151class SeriesTransformer(CustomTransformer):
1152 """Transformer for :class:`pandas.Series` objects."""
1154 @property
1155 def name(self) -> str:
1156 """Return transformer name."""
1157 return "series"
1159 def to_dict(self, transformer: "Transformer", o: object) -> dict[str, Any]:
1160 """Serialize a Series with its name, dtype, index, and data."""
1161 assert isinstance(o, pd.Series) # noqa: S101
1162 return {
1163 "name": transformer.to_dict(o.name),
1164 "dtype": str(o.dtype),
1165 "index": transformer.to_dict(o.index),
1166 "data": transformer.to_dict(o.tolist()),
1167 }
1169 def from_dict(self, transformer: "Transformer", d: dict[str, Any]) -> object:
1170 """Reconstruct a Series from its serialized form."""
1171 data = transformer.from_dict(d["data"])
1172 index = transformer.from_dict(d["index"])
1173 s = pd.Series(data, index=index, name=transformer.from_dict(d.get("name")))
1174 with contextlib.suppress(ValueError, TypeError):
1175 s = s.astype(d["dtype"])
1176 return s
1178 @property
1179 def supported_direct_types(self) -> Iterable[type]:
1180 """Return supported pandas Series type."""
1181 return [pd.Series]
1184class NodeKeyTransformer(CustomTransformer):
1185 """Transformer for :class:`~loman.nodekey.NodeKey` objects."""
1187 @property
1188 def name(self) -> str:
1189 """Return transformer name."""
1190 return "nodekey"
1192 def to_dict(self, transformer: "Transformer", o: object) -> dict[str, Any]:
1193 """Serialize a NodeKey as its path string."""
1194 return {"path": str(o)}
1196 def from_dict(self, transformer: "Transformer", d: dict[str, Any]) -> object:
1197 """Reconstruct a NodeKey from its path string."""
1198 from loman.nodekey import parse_nodekey
1200 return parse_nodekey(d["path"])
1202 @property
1203 def supported_direct_types(self) -> Iterable[type]:
1204 """Return supported NodeKey type."""
1205 from loman.nodekey import NodeKey
1207 return [NodeKey]