Coverage for src/loman/serialization/computation.py: 100%
337 statements
« prev ^ index » next coverage.py v7.16.0, created at 2026-09-07 00:32 +0000
1"""Serialization for Computation graphs to/from JSON."""
3from __future__ import annotations
5import contextvars
6import io
7import json
8import warnings
9import zipfile
10from pathlib import Path
11from typing import TYPE_CHECKING, Any, ClassVar, NamedTuple, TextIO
13from loman.consts import EdgeAttributes, NodeAttributes, States, SystemTags
14from loman.exception import SerializationError
15from loman.nodekey import parse_nodekey
17from .blobs import (
18 CONTAINER_STORE,
19 MANIFEST_NAME,
20 BlobReader,
21 BlobStore,
22 BlobWriter,
23 DirBlobStore,
24 ZipBlobStore,
25 read_dir_manifest,
26 read_zip_manifest,
27 write_dir_container,
28 write_zip_container,
29)
30from .profile import READABLE, SerializationProfile, resolve_profile
31from .transformer import (
32 DataFrameTransformer,
33 DefinitionMethodTransformer,
34 DillFunctionTransformer,
35 EnumTransformer,
36 FunctionRefTransformer,
37 NdArrayTransformer,
38 NodeKeyTransformer,
39 SeriesTransformer,
40 Transformer,
41 UntransformableTypeError,
42)
43from .values import register_value_transformers
45if TYPE_CHECKING:
46 pass
48# Serialization format version — bump when the schema changes.
49#
50# This is version 1: the first version of the format that is meant to be kept.
51# The shapes written by pre-release versions of loman are still read, on a
52# best-effort basis, by the reader branches marked "legacy" below --- an escaped
53# dict under "data" rather than "items", a row-major DataFrame, an index written
54# as a list of elements. Those predate any published guarantee and are not
55# themselves versioned, which is why the number restarts here rather than
56# continuing a sequence that never meant anything to anyone outside the repo.
57#
58# Note that the reader has never inspected this field, so a bump cannot make an
59# older release reject a newer file. Forward compatibility comes from keeping
60# changes additive, not from this number.
61FORMAT_VERSION = 1
63# Sentinel for a constant the policy dropped, distinct from a legitimate None.
64_DROPPED = object()
66_CONSTANT_POLICIES = frozenset({"raise", "drop"})
69class UnserializableConstantWarning(UserWarning):
70 """A constant argument was dropped, so the saved graph cannot be recalculated."""
73class UnserializableFunctionWarning(UserWarning):
74 """A node's function could not be encoded, so that node cannot be recalculated.
76 The node's *value* is still saved. What is lost is the ability to recompute
77 it: the reloaded node holds what it last produced and stays out of date
78 forever, because there is nothing left to call.
80 The common cause is a function that is not importable by name --- a lambda, a
81 closure, or a bound method of a class built by
82 :func:`~loman.computeengine.computation_factory`.
83 """
86CONTAINERS = ("zip", "dir", "json")
89class _SaveState(NamedTuple):
90 """The blob writer and profile in force for one save."""
92 writer: Any
93 profile: Any
96# Per-save state lives in a ContextVar rather than on the serializer, because a
97# serializer is a natural thing to build once and reuse --- and instance
98# attributes would mean two concurrent saves through the same one silently
99# writing into each other's archive. Measured: 11 of 12 threads failed that way.
100_SAVE_STATE: contextvars.ContextVar[_SaveState | None] = contextvars.ContextVar("loman_save_state", default=None)
102# Zip local file header. Every .loman archive starts with it.
103_ZIP_MAGIC = b"PK\x03\x04"
105# Pickle protocol 2+ opening opcode, as written by dill. Recognised only so a
106# write_dill file gets a useful error instead of "not valid JSON".
107_PICKLE_MAGIC = b"\x80"
110def infer_container_from_path(path: Path) -> str:
111 """Return the container implied by *path*'s suffix.
113 ``.json`` means the single readable document; everything else means a
114 ``.loman`` zip, which is the default because a single file is what people
115 move around.
116 """
117 return "json" if path.suffix.lower() == ".json" else "zip"
120def _resolve_profile_for_container(
121 profile: str | SerializationProfile | None,
122 container: str,
123 stores: dict[str, Any] | None = None,
124) -> SerializationProfile:
125 """Return the profile to use, rejecting the one impossible combination."""
126 if container not in CONTAINERS:
127 msg = f"Unknown container {container!r}; expected one of {list(CONTAINERS)}"
128 raise ValueError(msg)
130 if container == "json":
131 if profile is None:
132 return READABLE
133 resolved = resolve_profile(profile)
134 # A single JSON document holds no blobs of its own. It can still carry
135 # out-of-line values when an external store is supplied, so the refusal
136 # only applies when there is genuinely nowhere for them to go.
137 if resolved.inline_max_bytes is not None and not stores:
138 msg = (
139 f"The {resolved.name!r} profile writes values out of line, and a single JSON "
140 "document has nowhere to put them. Use container='zip' (or a .loman path) to "
141 "keep this profile, profile='readable' to keep the single document, or pass "
142 "stores={...} to hold the values somewhere of your own."
143 )
144 raise ValueError(msg)
145 return resolved
147 return resolve_profile(profile)
150def sniff_container(path: Path) -> str:
151 """Return which container *path* holds, by looking at it.
153 Order matters: a directory is checked first because it cannot be read as
154 bytes, then the zip magic number, then a leading brace for the single
155 document. A dill pickle is recognised only to say so, since "not valid JSON"
156 would be a poor description of that mistake.
157 """
158 if path.is_dir():
159 return "dir"
160 if not path.exists():
161 msg = f"No such file or directory: {str(path)!r}"
162 raise FileNotFoundError(msg)
164 with path.open("rb") as f:
165 head = f.read(4)
167 if head.startswith(_ZIP_MAGIC):
168 return "zip"
169 stripped = head.lstrip()
170 if stripped.startswith(b"{"):
171 return "json"
172 if head.startswith(_PICKLE_MAGIC):
173 msg = (
174 f"{str(path)!r} looks like a pickle written by write_dill, not a loman container. "
175 "Use Computation.read_dill to read it."
176 )
177 raise SerializationError(msg)
179 msg = (
180 f"Cannot tell what {str(path)!r} is. Expected a .loman archive, a directory "
181 f"containing {MANIFEST_NAME}, or a JSON document."
182 )
183 raise SerializationError(msg)
186def default_computation_transformer() -> Transformer:
187 """Create a Transformer pre-registered with all types needed for Computation serialization."""
188 t = Transformer()
190 # Numeric arrays
191 t.register(NdArrayTransformer())
193 # Enums: register the States enum so node states roundtrip correctly.
194 enum_t = EnumTransformer()
195 enum_t.register_enum(States)
196 t.register(enum_t)
198 # Importable callables (module-level functions). Lambdas / closures raise.
199 t.register(FunctionRefTransformer())
201 # Calc nodes declared on a @ComputationFactory class, which are bound methods
202 # of the definition object and so have no importable path of their own.
203 t.register(DefinitionMethodTransformer())
205 # Pandas
206 t.register(DataFrameTransformer())
207 t.register(SeriesTransformer())
209 # NodeKey (hierarchical node names)
210 t.register(NodeKeyTransformer())
212 # Dates, indexes, numpy scalars, sets, bytes, decimals.
213 register_value_transformers(t)
215 # Per-node execution timing. Imported here rather than at module scope
216 # because computeengine reaches back into this module to serialize.
217 from loman.computeengine import TimingData
219 t.register(TimingData)
221 return t
224def dill_computation_transformer() -> Transformer:
225 """Create a Transformer that serializes all callables — including lambdas and closures — via dill.
227 Identical to :func:`default_computation_transformer` except that
228 :class:`~loman.serialization.transformer.DillFunctionTransformer` replaces
229 :class:`~loman.serialization.transformer.FunctionRefTransformer`, so lambdas
230 and locally-defined closures are serialized as base64-encoded dill blobs
231 rather than raising :class:`~loman.exception.SerializationError`.
232 """
233 t = Transformer()
235 t.register(NdArrayTransformer())
237 enum_t = EnumTransformer()
238 enum_t.register_enum(States)
239 t.register(enum_t)
241 # Dill-based callable serializer — handles lambdas and closures.
242 t.register(DillFunctionTransformer())
244 t.register(DataFrameTransformer())
245 t.register(SeriesTransformer())
246 t.register(NodeKeyTransformer())
248 # Dates, indexes, numpy scalars, sets, bytes, decimals.
249 register_value_transformers(t)
251 # Per-node execution timing. Imported here rather than at module scope
252 # because computeengine reaches back into this module to serialize.
253 from loman.computeengine import TimingData
255 t.register(TimingData)
257 return t
260class ComputationSerializer:
261 """Serialize and deserialize a :class:`~loman.computeengine.Computation` graph to JSON.
263 The serialized format is a JSON object with the following top-level keys:
265 - ``version``: integer format version
266 - ``nodes``: list of node objects
267 - ``edges``: list of edge objects
269 Each **node** object has:
271 - ``key``: string representation of the NodeKey
272 - ``state``: name of the :class:`~loman.consts.States` enum member (or ``null``)
273 - ``value``: transformer-encoded value (or ``null`` when absent / not serialized)
274 - ``has_value``: bool — false when the node has no meaningful value to restore
275 - ``func``: transformer-encoded callable (or ``null``)
276 - ``args``: transformer-encoded constant positional arguments, keyed by
277 stringified positional index
278 - ``kwds``: transformer-encoded constant keyword arguments, keyed by parameter name
279 - ``serialize``: bool — whether the node has the ``__serialize__`` tag
280 - ``tags``: list of non-system tags
282 Each **edge** object has:
284 - ``src``: string key of the source node
285 - ``dst``: string key of the destination node
286 - ``param_type``: ``"arg"`` or ``"kwd"``
287 - ``param``: positional index (int) for args, parameter name (str) for kwds
289 Arguments taken from other nodes are recorded on **edges**, while arguments
290 given as :class:`~loman.computeengine.ConstantValue` are held on the node
291 itself and recorded in ``args`` and ``kwds``. Both are needed to call a node's
292 function, so a graph that dropped its constants would raise a
293 :class:`TypeError` the next time the node was calculated.
295 Parameters
296 ----------
297 transformer:
298 Custom :class:`~loman.serialization.transformer.Transformer` instance.
299 If ``None``, a default transformer is built based on *use_dill_for_functions*.
300 use_dill_for_functions:
301 When ``True``, lambdas and closures are serialized as base64-encoded dill
302 blobs rather than raising :class:`~loman.exception.SerializationError`.
303 Has no effect when a custom *transformer* is supplied. Defaults to ``False``.
304 on_unserializable_constant:
305 What to do when a constant argument cannot be encoded. ``"raise"``, the
306 default, refuses to write a graph that could not be recalculated.
307 ``"drop"`` omits the constant and emits
308 :class:`UnserializableConstantWarning`, restoring the behaviour of
309 releases before constants were recorded — where such a graph saved
310 silently and then raised :class:`TypeError` from the missing argument on
311 the first recalculation. It exists so an existing codebase can keep
312 writing files while it is fixed, not as a setting to leave in place.
313 """
315 # States that never carry a value worth writing. Everything else is
316 # serialized when the node actually holds a value --- notably STALE, whose
317 # value the in-memory computation keeps and whose intermediates are the
318 # whole point of saving a graph for post-mortem inspection.
319 _VALUELESS_STATES: ClassVar[set[States]] = {States.PLACEHOLDER, States.UNINITIALIZED}
321 def __init__(
322 self,
323 transformer: Transformer | None = None,
324 *,
325 use_dill_for_functions: bool = False,
326 on_unserializable_constant: str = "raise",
327 ) -> None:
328 """Initialise with an optional custom transformer."""
329 if on_unserializable_constant not in _CONSTANT_POLICIES:
330 msg = (
331 f"on_unserializable_constant must be one of {sorted(_CONSTANT_POLICIES)}, "
332 f"got {on_unserializable_constant!r}"
333 )
334 raise ValueError(msg)
335 if transformer is None:
336 transformer = (
337 dill_computation_transformer() if use_dill_for_functions else default_computation_transformer()
338 )
339 self._t = transformer
340 self._use_dill_for_functions = use_dill_for_functions
341 self._on_unserializable_constant = on_unserializable_constant
342 # Set per load() / loads() call; see the allow_code parameter there.
343 self._allow_code = True
345 def register(self, t: Any) -> None:
346 """Register a transformer or type with this serializer's transformer.
348 Accepts anything :meth:`~loman.serialization.transformer.Transformer.register`
349 accepts: a :class:`~loman.serialization.transformer.CustomTransformer`
350 instance, a :class:`~loman.serialization.transformer.Transformable`
351 subclass, an attrs class, or a dataclass.
353 The same serializer instance must be used for both writing and reading,
354 since the registration lives on the instance::
356 s = ComputationSerializer()
357 s.register(my_transformer)
358 comp.write_json('comp.json', serializer=s)
359 comp2 = Computation.read_json('comp.json', serializer=s)
360 """
361 self._t.register(t)
363 # ------------------------------------------------------------------
364 # Containers
365 # ------------------------------------------------------------------
367 def save(
368 self,
369 comp: Any,
370 path: str | Path,
371 *,
372 profile: str | SerializationProfile | None = None,
373 container: str | None = None,
374 stores: dict[str, BlobStore] | None = None,
375 ) -> None:
376 """Write *comp* to *path*.
378 :param path: Destination. A ``.json`` suffix implies the single-document
379 container; anything else defaults to a ``.loman`` zip.
380 :param profile: ``"readable"``, ``"efficient"``, or a
381 :class:`~loman.serialization.profile.SerializationProfile`. Defaults
382 to efficient, except in the ``json`` container where only readable is
383 possible.
384 :param container: ``"zip"``, ``"dir"`` or ``"json"``. Inferred from
385 *path* when omitted.
386 :param stores: Named :class:`~loman.serialization.blobs.BlobStore`
387 instances that nodes may be routed to. A node names a store through
388 ``add_node(store=...)`` or a profile override.
389 """
390 path = Path(path)
391 container = container or infer_container_from_path(path)
392 resolved = _resolve_profile_for_container(profile, container, stores)
393 external = dict(stores or {})
395 if container == "json":
396 with path.open("w", encoding="utf-8") as f:
397 self._dump_document(comp, f, resolved, external)
398 return
400 writer = write_zip_container if container == "zip" else write_dir_container
401 writer(
402 path,
403 lambda container_stores: self._build_manifest(comp, {**container_stores, **external}, resolved, container),
404 )
406 @staticmethod
407 def load_path(
408 path: str | Path,
409 *,
410 serializer: ComputationSerializer | None = None,
411 allow_code: bool = True,
412 stores: dict[str, BlobStore] | None = None,
413 ) -> Any:
414 """Read a computation from *path*, whatever container it uses.
416 :param stores: Named stores for blobs held outside the container. A
417 saved file records a store's name but never its configuration, so a
418 file with external blobs cannot resolve them unaided.
419 """
420 s = serializer if serializer is not None else ComputationSerializer()
421 path = Path(path)
422 container = sniff_container(path)
423 external = dict(stores or {})
425 if container == "json":
426 manifest = json.loads(path.read_text(encoding="utf-8"))
427 return s._read_manifest(manifest, external, allow_code=allow_code)
429 if container == "dir":
430 manifest = read_dir_manifest(path)
431 all_stores = {CONTAINER_STORE: DirBlobStore(path), **external}
432 return s._read_manifest(manifest, all_stores, allow_code=allow_code)
434 with zipfile.ZipFile(path) as zf:
435 manifest = read_zip_manifest(zf)
436 all_stores = {CONTAINER_STORE: ZipBlobStore(zf), **external}
437 return s._read_manifest(manifest, all_stores, allow_code=allow_code)
439 def _read_manifest(self, manifest: dict[str, Any], stores: dict[str, BlobStore], *, allow_code: bool) -> Any:
440 """Rebuild a computation from *manifest*, resolving blobs against *stores*."""
441 reader = BlobReader(manifest.get("blobs", []), stores)
442 with self._t.reading(reader):
443 return self._from_dict(manifest, allow_code=allow_code)
445 def _build_manifest(
446 self,
447 comp: Any,
448 stores: dict[str, BlobStore],
449 profile: SerializationProfile,
450 container: str,
451 ) -> dict[str, Any]:
452 """Return the manifest for *comp*, writing any blobs into *stores*."""
453 writer = BlobWriter(
454 stores,
455 compression=profile.compression,
456 dedupe=profile.dedupe,
457 checksums=profile.checksums,
458 )
459 token = _SAVE_STATE.set(_SaveState(writer=writer, profile=profile))
460 try:
461 manifest = self._to_dict(comp)
462 finally:
463 _SAVE_STATE.reset(token)
464 manifest["container"] = container
465 manifest["profile"] = profile.name
466 manifest["blobs"] = writer.table()
467 return manifest
469 def _dump_document(
470 self,
471 comp: Any,
472 fp: TextIO,
473 profile: SerializationProfile,
474 stores: dict[str, BlobStore] | None = None,
475 ) -> None:
476 """Write the single-document form.
478 The container itself holds no blobs, but an external store still can, so
479 a readable manifest can sit alongside data held elsewhere.
480 """
481 data = self._build_manifest(comp, dict(stores or {}), profile, "json")
482 json.dump(data, fp, allow_nan=False)
484 def dump(self, comp: Any, fp: TextIO) -> None:
485 """Serialize *comp* to *fp* (a text-mode file-like object)."""
486 self._dump_document(comp, fp, READABLE)
488 def dumps(self, comp: Any) -> str:
489 """Serialize *comp* and return a JSON string.
491 Always the readable single-document form: a string has nowhere to put
492 out-of-line bytes.
493 """
494 buf = io.StringIO()
495 self._dump_document(comp, buf, READABLE)
496 return buf.getvalue()
498 def _serialize_node_value(self, node_key: Any, state: States | None, node_data: dict[str, Any]) -> tuple[Any, bool]:
499 """Return ``(encoded_value, has_value)`` for a node that should be serialized.
501 Raises :class:`~loman.exception.SerializationError` if the value cannot
502 be encoded.
503 """
504 from loman.computeengine import Error
506 if state in self._VALUELESS_STATES or NodeAttributes.VALUE not in node_data:
507 return None, False
509 raw_value = node_data[NodeAttributes.VALUE]
510 # Keyed off the value's type, not the node's state. A node that failed
511 # and then went STALE --- because one of its inputs was replaced --- is no
512 # longer in ERROR state but still holds the Error it produced. Testing
513 # the state instead sent that value down the generic path, where an
514 # exception object has no encoding, and failed the entire save.
515 if isinstance(raw_value, Error):
516 exception_type = type(raw_value.exception)
517 return (
518 {
519 "__loman_error__": True,
520 "exception_type": exception_type.__name__,
521 "exception_module": exception_type.__module__,
522 "exception_str": str(raw_value.exception),
523 "traceback": raw_value.traceback,
524 },
525 True,
526 )
528 try:
529 return self._t.to_dict(raw_value), True
530 except (UntransformableTypeError, ValueError) as exc:
531 msg = f"Cannot serialize value of node {node_key!r}: {exc}"
532 raise SerializationError(msg) from exc
534 def _serialize_node_func(self, node_key: Any, raw_func: Any) -> Any:
535 """Return the encoded function for a node, or ``None`` if it cannot be serialized.
537 Lambdas raise :class:`~loman.exception.SerializationError` unless
538 ``use_dill_for_functions`` is enabled. Other non-importable callables
539 (e.g. framework closures from ``add_block``) are silently stored as ``null``.
540 """
541 qualname = getattr(raw_func, "__qualname__", "") or ""
542 if not self._use_dill_for_functions and "<lambda>" in qualname:
543 msg = (
544 f"Cannot serialize lambda function on node {node_key!r}. "
545 "Use a module-level importable function, serialize=False, "
546 "or ComputationSerializer(use_dill_for_functions=True)."
547 )
548 raise SerializationError(msg)
549 try:
550 return self._t.to_dict(raw_func)
551 except (UntransformableTypeError, ValueError, TypeError) as exc:
552 # The value is still saved; only the ability to recalculate is lost.
553 # That used to happen silently, so a graph could be reloaded, look
554 # complete, and never update again with nothing to explain why.
555 warnings.warn(
556 f"Cannot serialize the function on node {node_key!r} ({exc}). Its value is still "
557 "saved, but the reloaded node will have no function and so can never be "
558 "recalculated. Use a module-level importable function, or "
559 "ComputationSerializer(use_dill_for_functions=True).",
560 UnserializableFunctionWarning,
561 stacklevel=2,
562 )
563 return None
565 def _serialize_constant(self, node_key: Any, param: Any, value: Any) -> Any:
566 """Encode one constant argument held on a node.
568 Unlike a node function, a constant argument has no fallback: dropping it
569 would leave the node callable with the wrong number of arguments, so an
570 unrepresentable constant is an error rather than a ``null``.
571 """
572 try:
573 return self._t.to_dict(value)
574 except (UntransformableTypeError, ValueError, TypeError) as e:
575 msg = (
576 f"Cannot serialize constant argument {param!r} on node {node_key!r} ({e}). "
577 "Constant arguments are needed to call the node's function, so they cannot "
578 "be skipped: register a transformer for the type, set serialize=False on the "
579 "node, or use ComputationSerializer(use_dill_for_functions=True) for callables."
580 )
581 if self._on_unserializable_constant == "raise":
582 raise SerializationError(msg) from e
583 warnings.warn(
584 f"{msg} Dropping it, because on_unserializable_constant='drop'. The saved "
585 "graph will raise TypeError from the missing argument when this node is "
586 "recalculated.",
587 UnserializableConstantWarning,
588 stacklevel=2,
589 )
590 return _DROPPED
592 def _serialize_node_constants(
593 self, node_key: Any, node_data: dict[str, Any]
594 ) -> tuple[dict[str, Any], dict[str, Any]]:
595 """Return the encoded constant positional and keyword arguments for a node.
597 A constant the policy drops is omitted from the mapping entirely, so the
598 reloaded node looks exactly as it did before constants were recorded.
599 """
600 encoded_args = {
601 str(index): encoded
602 for index, value in node_data.get(NodeAttributes.ARGS, {}).items()
603 if (encoded := self._serialize_constant(node_key, index, value)) is not _DROPPED
604 }
605 encoded_kwds = {
606 name: encoded
607 for name, value in node_data.get(NodeAttributes.KWDS, {}).items()
608 if (encoded := self._serialize_constant(node_key, name, value)) is not _DROPPED
609 }
610 return encoded_args, encoded_kwds
612 def _serialize_node(self, node_key: Any, node_data: dict[str, Any]) -> dict[str, Any]:
613 """Return the serialized dict for a single node.
615 The whole node is encoded inside one write scope, so any value that asks
616 for out-of-line storage --- a node's value, or a constant argument that
617 happens to be a large array --- is attributed to this node in the blob
618 table.
619 """
620 state = _SAVE_STATE.get()
621 if state is None:
622 # save()/dumps() establish this; only reached if a caller drives
623 # _serialize_node directly.
624 return self._serialize_node_inner(node_key, node_data) # pragma: no cover
625 tags: set[str] = node_data.get(NodeAttributes.TAG, set())
626 with self._t.writing(
627 state.writer,
628 state.profile,
629 node=str(node_key),
630 tags=frozenset(tags),
631 store=node_data.get(NodeAttributes.STORE),
632 ):
633 return self._serialize_node_inner(node_key, node_data)
635 def _serialize_node_inner(self, node_key: Any, node_data: dict[str, Any]) -> dict[str, Any]:
636 """Return the serialized dict for a single node, inside a write scope."""
637 state: States | None = node_data.get(NodeAttributes.STATE)
638 tags: set[str] = node_data.get(NodeAttributes.TAG, set())
639 serialize_flag: bool = SystemTags.SERIALIZE in tags
641 serialized_state: States | None
642 if not serialize_flag:
643 serialized_state = States.UNINITIALIZED
644 encoded_value = None
645 has_value = False
646 else:
647 serialized_state = state
648 encoded_value, has_value = self._serialize_node_value(node_key, state, node_data)
650 raw_func = node_data.get(NodeAttributes.FUNC)
651 encoded_func = (
652 self._serialize_node_func(node_key, raw_func) if raw_func is not None and serialize_flag else None
653 )
654 encoded_args, encoded_kwds = (
655 self._serialize_node_constants(node_key, node_data) if encoded_func is not None else ({}, {})
656 )
658 user_tags = [t for t in tags if not t.startswith("__")]
660 out = {
661 "key": str(node_key),
662 "state": serialized_state.name if serialized_state is not None else None,
663 "value": encoded_value,
664 "has_value": has_value,
665 "func": encoded_func,
666 "args": encoded_args,
667 "kwds": encoded_kwds,
668 "serialize": serialize_flag,
669 "tags": user_tags,
670 }
671 out.update(self._serialize_node_attributes(node_key, node_data, serialize_flag=serialize_flag))
672 return out
674 def _serialize_node_attributes(
675 self, node_key: Any, node_data: dict[str, Any], *, serialize_flag: bool
676 ) -> dict[str, Any]:
677 """Return the presentational and execution attributes of a node.
679 Group, style and executor are plain strings describing how a node is
680 drawn and where it runs. They were previously dropped on load and
681 rebuilt as ``None``, so a reloaded graph rendered differently from the
682 one that was saved. The converter is a callable and follows the same
683 rules as the node's function; timing is written for the record and
684 restored as data.
685 """
686 attributes: dict[str, Any] = {}
687 for field, attr in (("group", NodeAttributes.GROUP), ("style", NodeAttributes.STYLE)):
688 value = node_data.get(attr)
689 if value is not None:
690 attributes[field] = self._t.to_dict(value)
692 executor = node_data.get(NodeAttributes.EXECUTOR)
693 if executor is not None:
694 attributes["executor"] = self._t.to_dict(executor)
696 store = node_data.get(NodeAttributes.STORE)
697 if store is not None:
698 attributes["store"] = self._t.to_dict(store)
700 converter = node_data.get(NodeAttributes.CONVERTER)
701 if converter is not None and serialize_flag:
702 encoded = self._serialize_node_func(node_key, converter)
703 if encoded is not None:
704 attributes["converter"] = encoded
706 timing = node_data.get(NodeAttributes.TIMING)
707 if timing is not None:
708 attributes["timing"] = self._t.to_dict(timing)
710 return attributes
712 def _restore_node_attributes(self, node_info: dict[str, Any], node_data: dict[str, Any]) -> None:
713 """Apply the presentational and execution attributes from *node_info*."""
714 node_data[NodeAttributes.GROUP] = self._t.from_dict(node_info.get("group"))
715 node_data[NodeAttributes.STYLE] = self._t.from_dict(node_info.get("style"))
716 node_data[NodeAttributes.EXECUTOR] = self._t.from_dict(node_info.get("executor"))
717 node_data[NodeAttributes.STORE] = self._t.from_dict(node_info.get("store"))
719 encoded_converter = node_info.get("converter")
720 node_data[NodeAttributes.CONVERTER] = self._decode_callable(encoded_converter)
722 encoded_timing = node_info.get("timing")
723 if encoded_timing is not None:
724 node_data[NodeAttributes.TIMING] = self._t.from_dict(encoded_timing)
726 def _serialize_edge(self, src: Any, dst: Any, edge_data: dict[str, Any]) -> dict[str, Any]:
727 """Return the serialized dict for a single edge."""
728 param = edge_data.get(EdgeAttributes.PARAM)
729 if param is None:
730 return {"src": str(src), "dst": str(dst), "param_type": None, "param": None}
732 from loman.computeengine import _ParameterType
734 param_type, param_val = param
735 return {
736 "src": str(src),
737 "dst": str(dst),
738 "param_type": "kwd" if param_type == _ParameterType.KWD else "arg",
739 "param": param_val,
740 }
742 def _to_dict(self, comp: Any) -> dict[str, Any]:
743 """Convert a Computation to a JSON-serializable dict."""
744 nodes_out = [self._serialize_node(k, comp.dag.nodes[k]) for k in comp.dag.nodes()]
745 edges_out = [self._serialize_edge(src, dst, data) for src, dst, data in comp.dag.edges(data=True)]
746 return {
747 "version": FORMAT_VERSION,
748 "nodes": nodes_out,
749 "edges": edges_out,
750 "metadata": self._serialize_metadata(comp),
751 }
753 def _serialize_metadata(self, comp: Any) -> dict[str, Any]:
754 """Return the computation's per-node metadata, keyed by node-key string.
756 Metadata is held on the computation rather than on dag nodes, so it needs
757 a map of its own. The root computation's own metadata is keyed by the
758 empty string.
759 """
760 from loman.nodekey import NodeKey
762 out: dict[str, Any] = {}
763 for node_key, metadata in getattr(comp, "_metadata", {}).items():
764 key = "" if node_key == NodeKey.root() else str(node_key)
765 out[key] = self._t.to_dict(metadata)
766 return out
768 def _restore_metadata(self, comp: Any, encoded: dict[str, Any]) -> None:
769 """Apply a serialized metadata map back onto *comp*."""
770 from loman.nodekey import NodeKey
772 for key, metadata in encoded.items():
773 node_key = NodeKey.root() if key == "" else parse_nodekey(key)
774 comp._metadata[node_key] = self._t.from_dict(metadata)
776 def load(self, fp: TextIO, *, allow_code: bool = True) -> Any:
777 """Deserialize a Computation from *fp* (a text-mode file-like object).
779 :param allow_code: When false, node functions and converters are not
780 restored. See :meth:`loads`.
781 """
782 data = json.load(fp)
783 return self._from_dict(data, allow_code=allow_code)
785 def loads(self, s: str, *, allow_code: bool = True) -> Any:
786 """Deserialize a Computation from a JSON string.
788 :param allow_code: When false, encoded callables are skipped rather than
789 resolved, and every node's function and converter comes back as
790 ``None``. Restoring a callable means importing the module the file
791 names, or unpickling a dill blob out of it --- both of which run code
792 the file chose. Values, structure, states and tags still load, which
793 is enough to inspect a graph from an untrusted source. Defaults to
794 true, preserving existing behaviour.
795 """
796 data = json.loads(s)
797 return self._from_dict(data, allow_code=allow_code)
799 def _decode_callable(self, encoded: Any) -> Any:
800 """Resolve an encoded callable, or ``None`` when code loading is refused."""
801 if encoded is None or not self._allow_code:
802 return None
803 return self._t.from_dict(encoded)
805 def _decode_error_value(self, encoded: dict[str, Any]) -> Any:
806 """Rebuild an ``Error`` from its recorded exception.
808 Builtin exception types are reconstructed so that ``except ValueError``
809 still matches after a round-trip. Anything else becomes a
810 :class:`~loman.exception.DeserializedError` carrying the original name:
811 importing the module a file names in order to rebuild its exception
812 would be running code chosen by the file.
813 """
814 import builtins
816 from loman.computeengine import Error
817 from loman.exception import DeserializedError
819 type_name = encoded.get("exception_type", "Exception")
820 module = encoded.get("exception_module")
821 message = encoded["exception_str"]
823 exception: Exception | None = None
824 if module in (None, "builtins"):
825 candidate = getattr(builtins, type_name, None)
826 # Exception, not BaseException: rebuilding a KeyboardInterrupt or a
827 # SystemExit as a node value would be a surprising thing to hand back.
828 if isinstance(candidate, type) and issubclass(candidate, Exception):
829 try:
830 exception = candidate(message)
831 except Exception: # pragma: no cover - exotic __init__
832 exception = None
833 if exception is None:
834 exception = DeserializedError(message, exception_type=type_name, exception_module=module)
836 return Error(exception=exception, traceback=encoded["traceback"])
838 def _from_dict(self, data: dict[str, Any], *, allow_code: bool = True) -> Any:
839 """Reconstruct a Computation from a deserialized dict."""
840 from loman.computeengine import Computation, _ParameterType
842 self._allow_code = allow_code
843 comp = Computation()
845 for node_info in data["nodes"]:
846 raw_key = node_info["key"]
847 node_key = parse_nodekey(raw_key)
848 state_name = node_info["state"]
849 state = States[state_name] if state_name is not None else None
850 serialize_flag: bool = node_info.get("serialize", True)
851 has_value: bool = node_info.get("has_value", False)
852 user_tags: list[str] = node_info.get("tags", [])
854 func = self._decode_callable(node_info.get("func"))
856 encoded_value = node_info.get("value")
857 if has_value:
858 # Decode even when the encoded value is null: a node whose value
859 # is legitimately None is not the same as a node with no value,
860 # and has_value is what distinguishes them.
861 if isinstance(encoded_value, dict) and encoded_value.get("__loman_error__"):
862 value = self._decode_error_value(encoded_value)
863 else:
864 value = self._t.from_dict(encoded_value)
865 else:
866 value = None
868 comp.dag.add_node(node_key)
869 node_data = comp.dag.nodes[node_key]
870 node_data[NodeAttributes.STATE] = state if state is not None else States.UNINITIALIZED
871 node_data[NodeAttributes.VALUE] = value if has_value else None
872 node_data[NodeAttributes.FUNC] = func
873 node_data[NodeAttributes.ARGS] = {
874 int(index): self._t.from_dict(encoded) for index, encoded in node_info.get("args", {}).items()
875 }
876 node_data[NodeAttributes.KWDS] = {
877 name: self._t.from_dict(encoded) for name, encoded in node_info.get("kwds", {}).items()
878 }
879 node_data[NodeAttributes.TAG] = set()
880 self._restore_node_attributes(node_info, node_data)
882 if serialize_flag:
883 node_data[NodeAttributes.TAG].add(SystemTags.SERIALIZE)
884 for tag in user_tags:
885 node_data[NodeAttributes.TAG].add(tag)
887 for edge_info in data["edges"]:
888 src_key = parse_nodekey(edge_info["src"])
889 dst_key = parse_nodekey(edge_info["dst"])
890 param_type_str = edge_info.get("param_type")
891 param_val = edge_info.get("param")
893 if param_type_str is not None:
894 param_type = _ParameterType.KWD if param_type_str == "kwd" else _ParameterType.ARG
895 comp.dag.add_edge(src_key, dst_key, **{EdgeAttributes.PARAM: (param_type, param_val)})
896 else:
897 comp.dag.add_edge(src_key, dst_key)
899 self._restore_metadata(comp, data.get("metadata", {}))
900 comp._refresh_maps()
902 return comp