Coverage for src/loman/computeengine.py: 99%
1240 statements
« prev ^ index » next coverage.py v7.16.0, created at 2026-09-07 00:32 +0000
1"""Core computation engine for dependency-aware calculation graphs."""
3import contextlib
4import functools
5import inspect
6import logging
7import traceback
8import types
9import warnings
10import weakref
11from collections import defaultdict
12from collections.abc import Callable, Hashable, Iterable, Mapping, Sequence
13from concurrent.futures import FIRST_COMPLETED, Executor, ThreadPoolExecutor, wait
14from dataclasses import dataclass, field
15from datetime import UTC, datetime
16from enum import Enum
17from types import MappingProxyType
18from typing import TYPE_CHECKING, Any, BinaryIO, TextIO, TypeVar, cast, overload
20if TYPE_CHECKING:
21 from .serialization.blobs import BlobStore
22 from .serialization.computation import ComputationSerializer
23 from .serialization.profile import SerializationProfile
24 from .ui import ComputationWidget
26import decorator
27import dill # nosec B403
28import networkx as nx
29import pandas as pd
31from .compat import get_signature
32from .consts import EdgeAttributes, NodeAttributes, NodeTransformations, States, SystemTags
33from .exception import (
34 CannotInsertToPlaceholderNodeException,
35 ComputationError,
36 LoopDetectedException,
37 MapException,
38 NodeAlreadyExistsException,
39 NonExistentNodeException,
40 ValidationError,
41)
42from .graph_utils import topological_sort
43from .nodekey import Name, Names, NodeKey, names_to_node_keys, node_keys_to_names, to_nodekey
44from .planning import ExecutionPlan, ValidationReport, create_execution_plan, validate_graph
45from .util import AttributeView, BlockFeature, RepeatedBlocks, apply1, apply_n, as_iterable, value_eq
46from .visualization import GraphView, NodeFormatter
48LOG = logging.getLogger("loman.computeengine")
50F = TypeVar("F", bound=Callable[..., Any])
53@dataclass
54class Error:
55 """Container for error information during computation."""
57 exception: Exception
58 traceback: str
61@dataclass
62class NodeData:
63 """Data associated with a computation node."""
65 state: States
66 value: object
69@dataclass
70class TimingData:
71 """Timing information for computation execution."""
73 start: datetime
74 end: datetime
75 duration: float
78@dataclass(frozen=True)
79class ComputationEvent:
80 """A batched notification describing a mutation to a computation.
82 Subscribers receive one event after each outermost public mutation, even
83 when that operation performs many internal state transitions. Values are
84 deliberately excluded: consumers can fetch a changed value lazily from
85 :attr:`computation` without copying large objects into every event.
87 :ivar computation: The live computation that produced the event. It is not a
88 snapshot, and continues to change after the event is delivered.
89 :ivar revision: Monotonic counter, matching :attr:`Computation.revision` at
90 the moment the event was published.
91 :ivar changed_nodes: Nodes whose state changed during the operation. When
92 :attr:`graph_changed` is true this is *not* a complete description of the
93 change, because adding, deleting or renaming nodes and altering tags or
94 styles need not change any node's state. Consumers reacting to a
95 structural event should re-read the graph rather than trusting this set.
96 :ivar states: The state of each entry in :attr:`changed_nodes` that still
97 exists, as of publication. Deleted nodes are absent.
98 :ivar graph_changed: True when the structure or presentation of the graph
99 changed, so any cached rendering of it is stale.
100 """
102 computation: "Computation"
103 revision: int
104 changed_nodes: frozenset[NodeKey]
105 states: Mapping[NodeKey, States]
106 graph_changed: bool = False
109ComputationSubscriber = Callable[[ComputationEvent], None]
111#: Cap on how many times subscriber-initiated mutations may cascade within one
112#: dispatch before Loman gives up. A well-behaved subscriber settles in one or
113#: two rounds; anything beyond this is a feedback loop rather than useful work.
114_MAX_NOTIFICATION_CASCADES = 16
117class _Subscription:
118 """One registered subscriber, held weakly when that is safe to do.
120 A callback bound to an object is held weakly, so subscribing a widget's
121 handler does not keep the widget alive for the lifetime of the computation.
122 Everything with no object behind it --- plain functions, lambdas, callable
123 objects, :func:`functools.partial` --- is held strongly, because callers
124 routinely pass a closure they retain no other reference to and holding
125 those weakly would collect them immediately.
127 "Bound to an object" means carrying a ``__self__``, which covers methods
128 written in Python and those written in C alike. The two need different
129 holders: :class:`weakref.WeakMethod` needs a ``__func__`` to rebind
130 against, and C methods have none, so those are held as a weak reference to
131 the owner plus the attribute name.
133 A few owners cannot be weakly referenced at all --- :class:`list`,
134 :class:`dict` and :class:`bytearray` among them --- so ``some_list.append``
135 falls back to a strong reference. That is a limitation of the type rather
136 than a decision here, and it errs towards a subscription that keeps
137 delivering rather than one that silently stops.
138 """
140 __slots__ = ("_name", "_owner", "_strong", "_weak")
142 def __init__(self, callback: ComputationSubscriber) -> None:
143 """Wrap ``callback``, choosing weak or strong ownership."""
144 self._weak: weakref.WeakMethod | None = None
145 self._owner: weakref.ref[Any] | None = None
146 self._name: str = ""
147 self._strong: ComputationSubscriber | None = None
148 if inspect.ismethod(callback):
149 self._weak = weakref.WeakMethod(callback)
150 return
151 owner = getattr(callback, "__self__", None)
152 name = getattr(callback, "__name__", None)
153 if owner is not None and name:
154 try:
155 self._owner = weakref.ref(owner)
156 except TypeError:
157 # list, dict, bytearray and friends support no weak references.
158 self._strong = callback
159 else:
160 self._name = name
161 return
162 self._strong = callback
164 def resolve(self) -> ComputationSubscriber | None:
165 """Return the callback, or ``None`` once a weakly held owner is gone."""
166 if self._strong is not None:
167 return self._strong
168 if self._weak is not None:
169 return self._weak()
170 if self._owner is None: # pragma: no cover - constructor covers all paths
171 return None
172 owner = self._owner()
173 return None if owner is None else getattr(owner, self._name, None)
176def _notifies_subscribers(*, graph_changed: bool = False) -> Callable[[F], F]:
177 """Batch changes made by a public mutation and notify on completion.
179 With no subscribers attached the wrapper is a straight pass-through, so
180 ordinary use of Loman pays nothing for the notification machinery.
181 """
183 def decorate(method: F) -> F:
184 @functools.wraps(method)
185 def wrapped(self: "Computation", *args: Any, **kwargs: Any) -> Any:
186 if self._change_depth == 0 and not self._subscriptions:
187 return method(self, *args, **kwargs)
188 self._change_depth += 1
189 try:
190 result = method(self, *args, **kwargs)
191 if graph_changed:
192 self._pending_graph_changed = True
193 return result
194 finally:
195 self._change_depth -= 1
196 if self._change_depth == 0:
197 self._publish_pending_events()
199 return cast("F", wrapped)
201 return decorate
204class _ParameterType(Enum):
205 """Internal enum for distinguishing positional and keyword parameters."""
207 ARG = 1
208 KWD = 2
211@dataclass
212class _ParameterItem:
213 """Internal container for parameter information during computation."""
215 type: _ParameterType
216 name: int | str
217 value: object
220def _node(func: Callable[..., Any], *args: Any, **kws: Any) -> Any: # pragma: no cover
221 """Internal wrapper function for node decorator."""
222 return func(*args, **kws)
225def node(comp: "Computation", name: Name | None = None, *args: Any, **kw: Any) -> Callable[[F], F]:
226 """Decorator to add a function as a node to a computation graph."""
228 def inner(f: F) -> F:
229 """Inner decorator that registers the function as a node."""
230 if name is None:
231 comp.add_node(f.__name__, f, *args, **kw)
232 else:
233 comp.add_node(name, f, *args, **kw)
234 result: F = decorator.decorate(f, _node)
235 return result
237 return inner
240@dataclass()
241class ConstantValue:
242 """Container for constant values in computations."""
244 value: object
247C = ConstantValue
250class Node:
251 """Base class for computation graph nodes."""
253 def add_to_comp(self, comp: "Computation", name: str, obj: object, ignore_self: bool) -> None:
254 """Add this node to the computation graph."""
255 raise NotImplementedError()
258@dataclass
259class InputNode(Node):
260 """A node representing input data in the computation graph."""
262 args: tuple[Any, ...] = field(default_factory=tuple)
263 kwds: dict[str, Any] = field(default_factory=dict)
265 def __init__(self, *args: Any, **kwds: Any) -> None:
266 """Initialize an input node with arguments and keyword arguments."""
267 self.args = args
268 self.kwds = kwds
270 def add_to_comp(self, comp: "Computation", name: str, obj: object, ignore_self: bool) -> None:
271 """Add this input node to the computation graph."""
272 comp.add_node(name, **self.kwds)
275input_node = InputNode
278def _bind_self(f: Any, obj: object, ignore_self: bool) -> Any:
279 """Bind a callback to the definition object when its first parameter is 'self'.
281 Anything that is not callable, including ``None`` and a plain node name, is
282 returned unchanged.
284 Asking for ``self`` when there is no definition object to bind to is a
285 contradiction, so it is reported here rather than as the bare
286 ``TypeError: instance must not be None`` that binding would otherwise raise.
287 """
288 if not callable(f) or not ignore_self:
289 return f
290 signature = get_signature(f)
291 if len(signature.kwd_params) > 0 and signature.kwd_params[0] == "self":
292 if obj is None:
293 name = getattr(f, "__qualname__", repr(f))
294 msg = (
295 f"Cannot bind 'self' for {name}: no definition object was supplied. "
296 "Pass one, drop the 'self' parameter, or use ignore_self=False."
297 )
298 raise ValueError(msg)
299 return types.MethodType(f, obj)
300 return f
303@dataclass
304class CalcNode(Node):
305 """A node representing a calculation in the computation graph."""
307 f: Callable[..., Any]
308 kwds: dict[str, Any] = field(default_factory=dict)
310 def add_to_comp(self, comp: "Computation", name: str, obj: object, ignore_self: bool) -> None:
311 """Add this calculation node to the computation graph."""
312 kwds = self.kwds.copy()
313 ignore_self = ignore_self or kwds.get("ignore_self", False)
314 f = self.f
315 if ignore_self:
316 signature = get_signature(self.f)
317 if len(signature.kwd_params) > 0 and signature.kwd_params[0] == "self":
318 f = f.__get__(obj, obj.__class__) # type: ignore[attr-defined]
319 if "ignore_self" in kwds:
320 del kwds["ignore_self"]
321 comp.add_node(name, f, **kwds)
324@overload
325def calc_node(f: F, **kwds: Any) -> F: ...
328@overload
329def calc_node(f: None = None, **kwds: Any) -> Callable[[F], F]: ...
332def calc_node(f: F | None = None, **kwds: Any) -> F | Callable[[F], F]:
333 """Decorator to mark a function as a calculation node."""
335 def wrap(func: F) -> F:
336 """Wrap function with node info attribute."""
337 func._loman_node_info = CalcNode(func, kwds)
338 return func
340 if f is None:
341 return wrap
342 return wrap(f)
345def _resolve_block(block: "Callable[[], Computation] | Computation") -> "Computation":
346 """Resolve a block definition, calling computation factories to build the block."""
347 if isinstance(block, Computation):
348 return block
349 if callable(block):
350 return block()
351 msg = f"Block {block} must be callable or Computation"
352 raise TypeError(msg)
355@dataclass
356class Block(Node):
357 """A node representing a computational block or subgraph."""
359 block: "Callable[[], Computation] | Computation"
360 args: tuple[Any, ...] = field(default_factory=tuple)
361 kwds: dict[str, Any] = field(default_factory=dict)
363 def __init__(self, block: "Callable[[], Computation] | Computation", *args: Any, **kwds: Any) -> None:
364 """Initialize a block node with a computation block and arguments."""
365 self.block = block
366 self.args = args
367 self.kwds = kwds
369 def add_to_comp(self, comp: "Computation", name: str, obj: object, ignore_self: bool) -> None:
370 """Add this block node to the computation graph."""
371 comp.add_block(name, _resolve_block(self.block), *self.args, **self.kwds)
374block = Block
377@dataclass
378class RepeatedBlocksNode(Node):
379 """A node representing one keyed copy of a computation block per key.
381 The attribute name used in a computation factory class becomes the base path
382 for the generated blocks, so ``instruments = repeated_blocks(...)`` with keys
383 ``('AAPL', 'MSFT')`` creates the blocks ``instruments/AAPL`` and
384 ``instruments/MSFT``. ``features`` describe how data flows in and out of every
385 copy, exactly as for :class:`loman.util.RepeatedBlocks`.
386 """
388 block: "Callable[[], Computation] | Computation"
389 keys: tuple[Hashable, ...] = field(default_factory=tuple)
390 features: tuple[BlockFeature, ...] = field(default_factory=tuple)
391 keep_values: bool = False
393 def __init__(
394 self,
395 block: "Callable[[], Computation] | Computation",
396 keys: Iterable[Hashable],
397 *,
398 features: Sequence[BlockFeature] = (),
399 keep_values: bool = False,
400 ) -> None:
401 """Initialize a repeated blocks node with a block template and its keys."""
402 self.block = block
403 self.keys = tuple(keys)
404 self.features = tuple(features)
405 self.keep_values = keep_values
407 def add_to_comp(self, comp: "Computation", name: str, obj: object, ignore_self: bool) -> None:
408 """Add the repeated blocks and the nodes their features describe."""
409 definition: RepeatedBlocks[Hashable] = RepeatedBlocks(
410 block=_resolve_block(self.block),
411 keys=self.keys,
412 base_path=name,
413 features=self.features,
414 keep_values=self.keep_values,
415 )
416 definition.add_to(comp, definition_object=obj, ignore_self=ignore_self)
419repeated_blocks = RepeatedBlocksNode
422def populate_computation_from_class(comp: "Computation", cls: type, obj: object, ignore_self: bool = True) -> None:
423 """Populate a computation from class methods with node decorators."""
424 for name, member in inspect.getmembers(cls):
425 node_: Node | None = None
426 if isinstance(member, Node):
427 node_ = member
428 elif hasattr(member, "_loman_node_info"):
429 node_ = member._loman_node_info
430 if node_ is not None:
431 node_.add_to_comp(comp, name, obj, ignore_self)
434def computation_factory(
435 maybe_cls: type | None = None, *, ignore_self: bool = True
436) -> Callable[..., "Computation"] | Callable[[type], Callable[..., "Computation"]]:
437 """Factory function to create computations from class definitions."""
439 def wrap(cls: type) -> Callable[..., "Computation"]:
440 """Wrap class to create computation factory function."""
442 @functools.wraps(cls, updated=())
443 def create_computation(*args: Any, **kwargs: Any) -> "Computation":
444 """Create a computation instance from the wrapped class."""
445 obj = cls()
446 comp = Computation(*args, **kwargs)
447 comp._definition_object = obj # type: ignore[attr-defined]
448 populate_computation_from_class(comp, cls, obj, ignore_self)
449 return comp
451 return create_computation
453 if maybe_cls is None:
454 return wrap
456 return wrap(maybe_cls)
459def _eval_node(
460 name: NodeKey,
461 f: Callable[..., Any],
462 args: list[Any],
463 kwds: dict[str, Any],
464 raise_exceptions: bool,
465) -> tuple[Any, Exception | None, str | None, datetime, datetime]:
466 """To make multiprocessing work, this function must be standalone so that pickle works."""
467 exc: Exception | None = None
468 tb: str | None = None
469 start_dt = datetime.now(UTC)
470 try:
471 logging.debug("Running " + str(name))
472 value = f(*args, **kwds)
473 logging.debug("Completed " + str(name))
474 except Exception as e:
475 value = None
476 exc = e
477 tb = traceback.format_exc()
478 if raise_exceptions:
479 raise
480 end_dt = datetime.now(UTC)
481 return value, exc, tb, start_dt, end_dt
484_MISSING_VALUE_SENTINEL = object()
487class NullObject:
488 """Debug helper object that raises exceptions for all attribute/item access."""
490 def __getattr__(self, name: str) -> Any:
491 """Raise AttributeError for any attribute access."""
492 print(f"__getattr__: {name}")
493 msg = f"'NullObject' object has no attribute '{name}'"
494 raise AttributeError(msg)
496 def __setattr__(self, name: str, value: Any) -> None:
497 """Raise AttributeError for any attribute assignment."""
498 print(f"__setattr__: {name}")
499 msg = f"'NullObject' object has no attribute '{name}'"
500 raise AttributeError(msg)
502 def __delattr__(self, name: str) -> None:
503 """Raise AttributeError for any attribute deletion."""
504 print(f"__delattr__: {name}")
505 msg = f"'NullObject' object has no attribute '{name}'"
506 raise AttributeError(msg)
508 def __call__(self, *args: Any, **kwargs: Any) -> Any:
509 """Raise TypeError when called as a function."""
510 print(f"__call__: {args}, {kwargs}")
511 msg = "'NullObject' object is not callable"
512 raise TypeError(msg)
514 def __getitem__(self, key: Any) -> Any:
515 """Raise KeyError for any item access."""
516 print(f"__getitem__: {key}")
517 msg = f"'NullObject' object has no item with key '{key}'"
518 raise KeyError(msg)
520 def __setitem__(self, key: Any, value: Any) -> None:
521 """Raise KeyError for any item assignment."""
522 print(f"__setitem__: {key}")
523 msg = f"'NullObject' object cannot have items set with key '{key}'"
524 raise KeyError(msg)
526 def __repr__(self) -> str:
527 """Return string representation of NullObject."""
528 print(f"__repr__: {object.__getattribute__(self, '__dict__')}")
529 return "<NullObject>"
532def identity_function(x: Any) -> Any:
533 """Return the input value unchanged."""
534 return x
537class Computation:
538 """A computation graph that manages dependencies and calculations.
540 The Computation class provides a framework for building and executing
541 computation graphs where nodes represent data or calculations, and edges
542 represent dependencies between them.
543 """
545 def __init__(
546 self,
547 *,
548 default_executor: Executor | None = None,
549 executor_map: dict[str, Executor] | None = None,
550 metadata: dict[str, Any] | None = None,
551 ) -> None:
552 """Initialize a new Computation.
554 :param default_executor: An executor
555 :type default_executor: concurrent.futures.Executor, default ThreadPoolExecutor(max_workers=1)
556 """
557 if default_executor is None:
558 self.default_executor: Executor = ThreadPoolExecutor(1)
559 else:
560 self.default_executor = default_executor
561 if executor_map is None:
562 self.executor_map: dict[str, Executor] = {}
563 else:
564 self.executor_map = executor_map
565 self.dag: nx.DiGraph = nx.DiGraph()
566 self._metadata: dict[NodeKey, Any] = {}
567 if metadata is not None:
568 self._metadata[NodeKey.root()] = metadata
570 self.v = self.get_attribute_view_for_path(NodeKey.root(), self._value_one, self.value)
571 self.s = self.get_attribute_view_for_path(NodeKey.root(), self._state_one, self.state)
572 self.i = self.get_attribute_view_for_path(NodeKey.root(), self._get_inputs_one_names, self.get_inputs)
573 self.o = self.get_attribute_view_for_path(NodeKey.root(), self._get_outputs_one, self.get_outputs)
574 self.t = self.get_attribute_view_for_path(NodeKey.root(), self._tag_one, self.tags)
575 self.style = self.get_attribute_view_for_path(NodeKey.root(), self._style_one, self.styles)
576 self.tim = self.get_attribute_view_for_path(NodeKey.root(), self._get_timing_one, self.get_timing)
577 self.x = self.get_attribute_view_for_path(
578 NodeKey.root(), self.compute_and_get_value, self.compute_and_get_value
579 )
580 self.src = self.get_attribute_view_for_path(NodeKey.root(), self.print_source, self.print_source)
581 self._tag_map: defaultdict[str, set[NodeKey]] = defaultdict(set)
582 self._state_map: dict[States, set[NodeKey]] = {state: set() for state in States}
583 self._subscriptions: list[_Subscription] = []
584 self._revision = 0
585 self._change_depth = 0
586 self._publishing = False
587 self._pending_changed_nodes: set[NodeKey] = set()
588 self._pending_graph_changed = False
590 @property
591 def revision(self) -> int:
592 """Return the revision number of the most recently published change."""
593 return self._revision
595 def subscribe(self, callback: ComputationSubscriber) -> Callable[[], None]:
596 """Subscribe to batched computation changes.
598 Subscribers are notified in registration order, synchronously, on the
599 thread that completes the outermost public mutation. A subscriber that
600 raises is logged and skipped; it never interrupts the mutation or the
601 other subscribers. A subscriber that itself mutates the computation
602 causes a further event to be published once the current round finishes,
603 up to a bounded number of cascades.
605 A callback with an object behind it --- anything carrying a
606 ``__self__``, whether written in Python or in C --- is held weakly, so
607 subscribing ``obj.handler`` or ``events.append`` does not keep the
608 owner alive; callers must retain it themselves. Everything else ---
609 plain functions, lambdas, callable objects, :func:`functools.partial`
610 --- is held strongly until unsubscribed, because callers commonly pass
611 a throwaway closure that nothing else references.
613 The exception is an owner that supports no weak references at all, such
614 as :class:`list`, :class:`dict` and :class:`bytearray`. There
615 ``some_list.append`` falls back to a strong reference, which is a
616 limitation of the type rather than a choice, and errs towards a
617 subscription that keeps delivering over one that silently stops.
619 Subscriptions are not copied by :meth:`copy` and are not serialized.
621 :param callback: Function accepting a :class:`ComputationEvent`.
622 :return: An idempotent, no-argument unsubscribe function.
623 """
624 if not callable(callback):
625 msg = "callback must be callable"
626 raise TypeError(msg)
627 subscription = _Subscription(callback)
628 self._subscriptions.append(subscription)
630 def unsubscribe() -> None:
631 """Remove this subscription, ignoring repeat calls."""
632 with contextlib.suppress(ValueError):
633 self._subscriptions.remove(subscription)
635 return unsubscribe
637 def _mark_changed(self, *node_keys: NodeKey) -> None:
638 """Record nodes changed by the current public mutation."""
639 if self._subscriptions:
640 self._pending_changed_nodes.update(node_keys)
642 def _take_pending_event(self) -> ComputationEvent:
643 """Consume the batched changes and turn them into one event."""
644 changed_nodes = frozenset(self._pending_changed_nodes)
645 graph_changed = self._pending_graph_changed
646 self._pending_changed_nodes.clear()
647 self._pending_graph_changed = False
648 self._revision += 1
649 states = {
650 node_key: self.dag.nodes[node_key][NodeAttributes.STATE]
651 for node_key in changed_nodes
652 if node_key in self.dag
653 }
654 return ComputationEvent(self, self._revision, changed_nodes, MappingProxyType(states), graph_changed)
656 def _publish_pending_events(self) -> None:
657 """Publish batched events, including any a subscriber triggers in turn.
659 Re-entrant calls return immediately: the dispatch loop already running
660 picks up whatever the subscriber changed, so a subscriber that mutates
661 the computation cannot recurse into the stack.
662 """
663 if self._publishing:
664 return
665 self._publishing = True
666 try:
667 for _ in range(_MAX_NOTIFICATION_CASCADES):
668 if not self._pending_changed_nodes and not self._pending_graph_changed:
669 return
670 self._dispatch(self._take_pending_event())
671 if self._pending_changed_nodes or self._pending_graph_changed:
672 LOG.error(
673 "Computation subscribers kept mutating the computation after %s rounds; "
674 "discarding further notifications to break the loop",
675 _MAX_NOTIFICATION_CASCADES,
676 )
677 self._pending_changed_nodes.clear()
678 self._pending_graph_changed = False
679 finally:
680 self._publishing = False
682 def _dispatch(self, event: ComputationEvent) -> None:
683 """Deliver one event to every live subscriber, isolating failures."""
684 dead = False
685 for subscription in tuple(self._subscriptions):
686 callback = subscription.resolve()
687 if callback is None:
688 dead = True
689 continue
690 try:
691 callback(event)
692 except Exception:
693 LOG.exception("Computation subscriber failed at revision %s", event.revision)
694 if dead:
695 self._subscriptions = [s for s in self._subscriptions if s.resolve() is not None]
697 def get_attribute_view_for_path(
698 self, nodekey: NodeKey, get_one_func: Callable[[Name], Any], get_many_func: Callable[[Name | Names], Any]
699 ) -> AttributeView:
700 """Create an attribute view for a specific node path."""
702 def node_func() -> Iterable[str]:
703 """Return list of child node names for this path."""
704 return [str(n) for n in self.get_tree_list_children(nodekey)]
706 def get_one_func_for_path(name: str) -> Any:
707 """Get value for a single node at this path."""
708 nk = to_nodekey(name)
709 new_nk = nk.prepend(nodekey)
710 if self.has_node(new_nk):
711 return get_one_func(new_nk)
712 elif self.tree_has_path(new_nk):
713 return self.get_attribute_view_for_path(new_nk, get_one_func, get_many_func)
714 else:
715 msg = f"Path {new_nk} does not exist"
716 raise KeyError(msg) # pragma: no cover
718 def get_many_func_for_path(name: Name | Names) -> Any:
719 """Get values for one or more nodes at this path."""
720 if isinstance(name, list):
721 return [get_one_func_for_path(str(n)) for n in name]
722 else:
723 return get_one_func_for_path(str(name))
725 return AttributeView(node_func, get_one_func_for_path, get_many_func_for_path)
727 def _get_names_for_state(self, state: States) -> set[Name]:
728 """Get node names that have a specific state."""
729 return set(node_keys_to_names(self._state_map[state]))
731 def _get_tags_for_state(self, tag: str) -> set[Name]:
732 """Get node names that have a specific tag."""
733 return set(node_keys_to_names(self._tag_map[tag]))
735 def _process_function_args(self, node_key: NodeKey, node: dict[str, Any], args: list[Any] | None) -> int:
736 """Process positional arguments for a function node."""
737 args_count = 0
738 if args:
739 args_count = len(args)
740 for i, arg in enumerate(args):
741 if isinstance(arg, ConstantValue):
742 node[NodeAttributes.ARGS][i] = arg.value
743 else:
744 input_vertex_name = arg
745 input_vertex_node_key = to_nodekey(input_vertex_name)
746 if not self.dag.has_node(input_vertex_node_key):
747 self.dag.add_node(input_vertex_node_key, **{NodeAttributes.STATE: States.PLACEHOLDER})
748 self._state_map[States.PLACEHOLDER].add(input_vertex_node_key)
749 self.dag.add_edge(
750 input_vertex_node_key, node_key, **{EdgeAttributes.PARAM: (_ParameterType.ARG, i)}
751 )
752 return args_count
754 def _build_param_map(
755 self,
756 func: Callable[..., Any],
757 node_key: NodeKey,
758 args_count: int,
759 kwds: dict[str, Any] | None,
760 inspect: bool,
761 ) -> tuple[dict[str, Any], list[str]]:
762 """Build parameter map for function node."""
763 param_map: dict[str, Any] = {}
764 default_names: list[str] = []
766 if inspect:
767 signature = get_signature(func)
768 if not signature.has_var_args:
769 for param_name in signature.kwd_params[args_count:]:
770 if kwds is not None and param_name in kwds:
771 param_source = kwds[param_name]
772 else:
773 param_source = node_key.parent.join_parts(param_name)
774 param_map[param_name] = param_source
775 if signature.has_var_kwds and kwds is not None:
776 for param_name, param_source in kwds.items():
777 param_map[param_name] = param_source
778 default_names = signature.default_params
779 else:
780 if kwds is not None:
781 for param_name, param_source in kwds.items():
782 param_map[param_name] = param_source
784 return param_map, default_names
786 def _process_function_kwds(
787 self, node_key: NodeKey, node: dict[str, Any], param_map: dict[str, Any], default_names: list[str]
788 ) -> None:
789 """Process keyword arguments for a function node."""
790 for param_name, param_source in param_map.items():
791 if isinstance(param_source, ConstantValue):
792 node[NodeAttributes.KWDS][param_name] = param_source.value
793 else:
794 in_node_name = param_source
795 in_node_key = to_nodekey(in_node_name)
796 if not self.dag.has_node(in_node_key):
797 if param_name in default_names:
798 continue
799 else:
800 self.dag.add_node(in_node_key, **{NodeAttributes.STATE: States.PLACEHOLDER})
801 self._state_map[States.PLACEHOLDER].add(in_node_key)
802 self.dag.add_edge(in_node_key, node_key, **{EdgeAttributes.PARAM: (_ParameterType.KWD, param_name)})
804 @_notifies_subscribers(graph_changed=True)
805 def add_node(
806 self,
807 name: Name,
808 func: Callable[..., Any] | None = None,
809 *,
810 args: list[Any] | None = None,
811 kwds: dict[str, Any] | None = None,
812 value: Any = _MISSING_VALUE_SENTINEL,
813 converter: Callable[[Any], Any] | None = None,
814 serialize: bool = True,
815 inspect: bool = True,
816 group: str | None = None,
817 tags: Iterable[str] | None = None,
818 style: str | None = None,
819 executor: str | None = None,
820 store: str | None = None,
821 metadata: dict[str, Any] | None = None,
822 ) -> None:
823 """Adds or updates a node in a computation.
825 :param name: Name of the node to add. This may be any hashable object.
826 :param func: Function to use to calculate the node if the node is a calculation node. By default, the input
827 nodes to the function will be implied from the names of the function parameters. For example, a
828 parameter called ``a`` would be taken from the node called ``a``. This can be modified with the
829 ``kwds`` parameter.
830 :type func: Function, default None
831 :param args: Specifies a list of nodes that will be used to populate arguments of the function positionally
832 for a calculation node. e.g. If args is ``['a', 'b', 'c']`` then the function would be called with
833 three parameters, taken from the nodes 'a', 'b' and 'c' respectively.
834 :type args: List, default None
835 :param kwds: Specifies a mapping from parameter name to the node that should be used to populate that
836 parameter when calling the function for a calculation node. e.g. If args is ``{'x': 'a', 'y': 'b'}``
837 then the function would be called with parameters named 'x' and 'y', and their values would be taken
838 from nodes 'a' and 'b' respectively. Each entry in the dictionary can be read as "take parameter
839 [key] from node [value]".
840 :type kwds: Dictionary, default None
841 :param value: If given, the value is inserted into the node, and the node state set to UPTODATE.
842 :type value: default None
843 :param converter: Callable applied to any value on its way into the node. The node stores what the
844 converter returns, both for values supplied by ``value``, ``insert`` and ``insert_many``, and for
845 values the node calculates with ``func``. A converter that raises leaves the node in state ERROR
846 without storing the value, which is how a validator is written: check the value and return it
847 unchanged when it is acceptable. The exception propagates to the caller when the value was
848 supplied, but not when it was calculated, where the failure is reported as node state instead.
849 A converter is saved by reference, like a node's function, so it must be importable: a
850 module-level function or builtin round-trips, while a lambda raises ``SerializationError``.
851 :type converter: Callable, default None
852 :param serialize: Whether the node should be serialized. Some objects cannot be serialized, in which
853 case, set serialize to False
854 :type serialize: boolean, default True
855 :param inspect: Whether to use introspection to determine the arguments of the function, which can be
856 slow. If this is not set, kwds and args must be set for the function to obtain parameters.
857 :type inspect: boolean, default True
858 :param group: Subgraph to render node in
859 :type group: default None
860 :param tags: Set of tags to apply to node
861 :type tags: Iterable
862 :param styles: Style to apply to node
863 :type styles: String, default None
864 :param executor: Name of executor to run node on
865 :type executor: string
866 :param store: Name of the blob store this node's value should be saved
867 to, for values that belong somewhere other than the saved file --- a
868 bucket, a database. The store itself is supplied at save and load
869 time as ``stores={name: ...}``, so the graph names a destination
870 without holding a bucket or a credential. A profile override for the
871 same node wins over this, which is what lets one computation be
872 saved to that store in production and to a plain container in a test.
873 :type store: string, default None
874 """
875 node_key = to_nodekey(name)
876 LOG.debug(f"Adding node {node_key}")
877 has_value = value is not _MISSING_VALUE_SENTINEL
878 if value is _MISSING_VALUE_SENTINEL:
879 value = None
880 if tags is None:
881 tags = []
883 self.dag.add_node(node_key)
884 pred_edges = [(p, node_key) for p in self.dag.predecessors(node_key)]
885 self.dag.remove_edges_from(pred_edges)
886 node = self.dag.nodes[node_key]
888 if metadata is None:
889 if node_key in self._metadata:
890 del self._metadata[node_key]
891 else:
892 self._metadata[node_key] = metadata
894 self._set_state_and_literal_value(node_key, States.UNINITIALIZED, None, require_old_state=False)
896 node[NodeAttributes.TAG] = set()
897 node[NodeAttributes.STYLE] = style
898 node[NodeAttributes.GROUP] = group
899 node[NodeAttributes.ARGS] = {}
900 node[NodeAttributes.KWDS] = {}
901 node[NodeAttributes.FUNC] = None
902 node[NodeAttributes.EXECUTOR] = executor
903 node[NodeAttributes.CONVERTER] = converter
904 node[NodeAttributes.STORE] = store
906 if func:
907 node[NodeAttributes.FUNC] = func
908 args_count = self._process_function_args(node_key, node, args)
909 param_map, default_names = self._build_param_map(func, node_key, args_count, kwds, inspect)
910 self._process_function_kwds(node_key, node, param_map, default_names)
911 self._set_descendents(node_key, States.STALE)
913 if has_value:
914 self._set_uptodate(node_key, value)
915 if node[NodeAttributes.STATE] == States.UNINITIALIZED:
916 self._try_set_computable(node_key)
917 self.set_tag(node_key, tags)
918 if serialize:
919 self.set_tag(node_key, SystemTags.SERIALIZE)
921 def _refresh_maps(self) -> None:
922 """Refresh internal tag and state maps from node data."""
923 self._tag_map.clear()
924 for state in States:
925 self._state_map[state].clear()
926 for node_key in self._node_keys():
927 state = self.dag.nodes[node_key][NodeAttributes.STATE]
928 self._state_map[state].add(node_key)
929 tags = self.dag.nodes[node_key].get(NodeAttributes.TAG, set())
930 for tag in tags:
931 self._tag_map[tag].add(node_key)
933 def _set_tag_one(self, name: Name, tag: str) -> None:
934 """Set a single tag on a single node."""
935 node_key = to_nodekey(name)
936 self.dag.nodes[node_key][NodeAttributes.TAG].add(tag)
937 self._tag_map[tag].add(node_key)
939 @_notifies_subscribers(graph_changed=True)
940 def set_tag(self, name: Name | Names, tag: str | Iterable[str]) -> None:
941 """Set tags on a node or nodes. Ignored if tags are already set.
943 :param name: Node or nodes to set tag for
944 :param tag: Tag to set
945 """
946 apply_n(self._set_tag_one, name, tag)
948 def _clear_tag_one(self, name: Name, tag: str) -> None:
949 """Clear a single tag from a single node."""
950 node_key = to_nodekey(name)
951 self.dag.nodes[node_key][NodeAttributes.TAG].discard(tag)
952 self._tag_map[tag].discard(node_key)
954 @_notifies_subscribers(graph_changed=True)
955 def clear_tag(self, name: Name | Names, tag: str | Iterable[str]) -> None:
956 """Clear tag on a node or nodes. Ignored if tags are not set.
958 :param name: Node or nodes to clear tags for
959 :param tag: Tag to clear
960 """
961 apply_n(self._clear_tag_one, name, tag)
963 def _set_style_one(self, name: Name, style: str) -> None:
964 """Set style on a single node."""
965 node_key = to_nodekey(name)
966 self.dag.nodes[node_key][NodeAttributes.STYLE] = style
968 @_notifies_subscribers(graph_changed=True)
969 def set_style(self, name: Name | Names, style: str) -> None:
970 """Set styles on a node or nodes.
972 :param name: Node or nodes to set style for
973 :param style: Style to set
974 """
975 apply_n(self._set_style_one, name, style)
977 def _clear_style_one(self, name: Name) -> None:
978 """Clear style from a single node."""
979 node_key = to_nodekey(name)
980 self.dag.nodes[node_key][NodeAttributes.STYLE] = None
982 @_notifies_subscribers(graph_changed=True)
983 def clear_style(self, name: Name | Names) -> None:
984 """Clear style on a node or nodes.
986 :param name: Node or nodes to clear styles for
987 """
988 apply_n(self._clear_style_one, name)
990 def metadata(self, name: Name) -> dict[str, Any]:
991 """Get metadata for a node."""
992 node_key = to_nodekey(name)
993 if self.tree_has_path(name):
994 if node_key not in self._metadata:
995 self._metadata[node_key] = {}
996 result: dict[str, Any] = self._metadata[node_key]
997 return result
998 else:
999 msg = f"Node {node_key} does not exist."
1000 raise NonExistentNodeException(msg)
1002 @_notifies_subscribers(graph_changed=True)
1003 def delete_node(self, name: Name) -> None:
1004 """Delete a node from a computation.
1006 When nodes are explicitly deleted with ``delete_node``, but are still depended on by other nodes, then they
1007 will be set to PLACEHOLDER status. In this case, if the nodes that depend on a PLACEHOLDER node are deleted,
1008 then the PLACEHOLDER node will also be deleted.
1010 :param name: Name of the node to delete. If the node does not exist, a ``NonExistentNodeException`` will
1011 be raised.
1012 """
1013 node_key = to_nodekey(name)
1014 LOG.debug(f"Deleting node {node_key}")
1016 if not self.dag.has_node(node_key):
1017 msg = f"Node {node_key} does not exist"
1018 raise NonExistentNodeException(msg)
1020 if node_key in self._metadata:
1021 del self._metadata[node_key]
1023 if len(self.dag.succ[node_key]) == 0:
1024 preds = self.dag.predecessors(node_key)
1025 state = self.dag.nodes[node_key][NodeAttributes.STATE]
1026 self.dag.remove_node(node_key)
1027 self._state_map[state].remove(node_key)
1028 self._mark_changed(node_key)
1029 for n in preds:
1030 if self.dag.nodes[n][NodeAttributes.STATE] == States.PLACEHOLDER:
1031 self.delete_node(n)
1032 else:
1033 self._set_state(node_key, States.PLACEHOLDER)
1035 @_notifies_subscribers(graph_changed=True)
1036 def rename_node(self, old_name: Name | Mapping[Name, Name], new_name: Name | None = None) -> None:
1037 """Rename a node in a computation.
1039 :param old_name: Node to rename, or a dictionary of nodes to rename, with existing names as keys, and
1040 new names as values
1041 :param new_name: New name for node.
1042 """
1043 name_mapping: dict[Name, Name]
1044 if isinstance(old_name, Mapping) and not isinstance(old_name, str):
1045 for k, v in old_name.items():
1046 LOG.debug(f"Renaming node {k} to {v}")
1047 if new_name is not None:
1048 msg = "new_name must not be set if rename_node is passed a dictionary"
1049 raise ValueError(msg)
1050 else:
1051 name_mapping = dict(old_name) # type: ignore[arg-type]
1052 else:
1053 LOG.debug(f"Renaming node {old_name} to {new_name}")
1054 old_node_key = to_nodekey(old_name)
1055 if not self.dag.has_node(old_node_key):
1056 msg = f"Node {old_name} does not exist"
1057 raise NonExistentNodeException(msg)
1058 assert new_name is not None # noqa: S101
1059 new_node_key = to_nodekey(new_name)
1060 if self.dag.has_node(new_node_key):
1061 msg = f"Node {new_name} already exists"
1062 raise NodeAlreadyExistsException(msg)
1063 name_mapping = {old_name: new_name}
1065 node_key_mapping = {to_nodekey(on): to_nodekey(nn) for on, nn in name_mapping.items()}
1066 nx.relabel_nodes(self.dag, node_key_mapping, copy=False)
1068 for old_nk, new_nk in node_key_mapping.items():
1069 if old_nk in self._metadata:
1070 self._metadata[new_nk] = self._metadata[old_nk]
1071 del self._metadata[old_nk]
1072 else:
1073 if new_nk in self._metadata: # pragma: no cover
1074 del self._metadata[new_nk]
1076 self._mark_changed(*node_key_mapping.keys(), *node_key_mapping.values())
1077 self._refresh_maps()
1079 @_notifies_subscribers(graph_changed=True)
1080 def repoint(self, old_name: Name, new_name: Name) -> None:
1081 """Changes all nodes that use old_name as an input to use new_name instead.
1083 Note that if old_name is an input to new_name, then that will not be changed, to try to avoid introducing
1084 circular dependencies, but other circular dependencies will not be checked.
1086 If new_name does not exist, then it will be created as a PLACEHOLDER node.
1088 :param old_name:
1089 :param new_name:
1090 :return:
1091 """
1092 old_node_key = to_nodekey(old_name)
1093 new_node_key = to_nodekey(new_name)
1094 if old_node_key == new_node_key:
1095 return
1097 changed_names = list(self.dag.successors(old_node_key))
1099 if len(changed_names) > 0 and not self.dag.has_node(new_node_key):
1100 self.dag.add_node(new_node_key, **{NodeAttributes.STATE: States.PLACEHOLDER})
1101 self._state_map[States.PLACEHOLDER].add(new_node_key)
1103 for name in changed_names:
1104 if name == new_node_key:
1105 continue
1106 edge_data = self.dag.get_edge_data(old_node_key, name)
1107 self.dag.add_edge(new_node_key, name, **edge_data)
1108 self.dag.remove_edge(old_node_key, name)
1110 for n in changed_names:
1111 self.set_stale(n)
1113 @_notifies_subscribers()
1114 def insert(self, name: Name, value: Any, force: bool = False) -> None:
1115 """Insert a value into a node of a computation.
1117 Following insertation, the node will have state UPTODATE, and all its descendents will be COMPUTABLE or STALE.
1119 If an attempt is made to insert a value into a node that does not exist, a ``NonExistentNodeException``
1120 will be raised.
1122 :param name: Name of the node to add.
1123 :param value: The value to be inserted into the node.
1124 :param force: Whether to force recalculation of descendents if node value and state would not be changed
1125 """
1126 node_key = to_nodekey(name)
1127 LOG.debug(f"Inserting value into node {node_key}")
1129 if not self.dag.has_node(node_key):
1130 msg = f"Node {node_key} does not exist"
1131 raise NonExistentNodeException(msg)
1133 state = self._state_one(name)
1134 if state == States.PLACEHOLDER:
1135 msg = "Cannot insert into placeholder node. Use add_node to create the node first"
1136 raise CannotInsertToPlaceholderNodeException(msg)
1138 if not force and state == States.UPTODATE:
1139 current_value = self._value_one(name)
1140 if value_eq(value, current_value):
1141 return
1143 self._set_state_and_value(node_key, States.UPTODATE, value)
1144 self._set_descendents(node_key, States.STALE)
1145 for n in self.dag.successors(node_key):
1146 self._try_set_computable(n)
1148 @_notifies_subscribers()
1149 def insert_many(self, name_value_pairs: Iterable[tuple[Name, object]]) -> None:
1150 """Insert values into many nodes of a computation simultaneously.
1152 Following insertation, the nodes will have state UPTODATE, and all their descendents will be COMPUTABLE
1153 or STALE. In the case of inserting many nodes, some of which are descendents of others, this ensures that
1154 the inserted nodes have correct status, rather than being set as STALE when their ancestors are inserted.
1156 If an attempt is made to insert a value into a node that does not exist, a ``NonExistentNodeException`` will be
1157 raised, and none of the nodes will be inserted.
1159 :param name_value_pairs: Each tuple should be a pair (name, value), where name is the name of the node to
1160 insert the value into.
1161 :type name_value_pairs: List of tuples
1162 """
1163 node_key_value_pairs = [(to_nodekey(name), value) for name, value in name_value_pairs]
1164 LOG.debug(f"Inserting value into nodes {', '.join(str(name) for name, value in node_key_value_pairs)}")
1166 for name, _value in node_key_value_pairs:
1167 if not self.dag.has_node(name):
1168 msg = f"Node {name} does not exist"
1169 raise NonExistentNodeException(msg)
1171 stale = set()
1172 computable = set()
1173 for name, value in node_key_value_pairs:
1174 self._set_state_and_value(name, States.UPTODATE, value)
1175 stale.update(nx.dag.descendants(self.dag, name))
1176 computable.update(self.dag.successors(name))
1177 names = {name for name, value in node_key_value_pairs}
1178 stale.difference_update(names)
1179 computable.difference_update(names)
1180 for name in stale:
1181 self._set_state(name, States.STALE)
1182 for name in computable:
1183 self._try_set_computable(name)
1185 @_notifies_subscribers()
1186 def insert_from(self, other: "Computation", nodes: Iterable[Name] | None = None) -> None:
1187 """Insert values into another Computation object into this Computation object.
1189 :param other: The computation object to take values from
1190 :type Computation:
1191 :param nodes: Only populate the nodes with the names provided in this list. By default, all nodes from the
1192 other Computation object that have corresponding nodes in this Computation object will be inserted
1193 :type nodes: List, default None
1194 """
1195 if nodes is None:
1196 nodes_set: set[Any] = set(self.dag.nodes)
1197 nodes_set.intersection_update(other.dag.nodes())
1198 nodes = nodes_set
1199 name_value_pairs = [(name, other.value(name)) for name in nodes]
1200 self.insert_many(name_value_pairs)
1202 def _set_state(self, node_key: NodeKey, state: States) -> None:
1203 """Set the state of a node without changing its value."""
1204 node = self.dag.nodes[node_key]
1205 old_state = node[NodeAttributes.STATE]
1206 self._state_map[old_state].remove(node_key)
1207 node[NodeAttributes.STATE] = state
1208 self._state_map[state].add(node_key)
1209 self._mark_changed(node_key)
1211 def _set_state_and_value(
1212 self, node_key: NodeKey, state: States, value: object, *, throw_conversion_exception: bool = True
1213 ) -> None:
1214 """Set state and value of a node, applying any converter."""
1215 node = self.dag.nodes[node_key]
1216 converter = node.get(NodeAttributes.CONVERTER)
1217 if converter is None:
1218 self._set_state_and_literal_value(node_key, state, value)
1219 else:
1220 try:
1221 converted_value = converter(value)
1222 self._set_state_and_literal_value(node_key, state, converted_value)
1223 except Exception as e:
1224 tb = traceback.format_exc()
1225 self._set_error(node_key, e, tb)
1226 if throw_conversion_exception:
1227 raise
1229 def _set_state_and_literal_value(
1230 self, node_key: NodeKey, state: States, value: object, require_old_state: bool = True
1231 ) -> None:
1232 """Set state and literal value of a node without conversion."""
1233 node = self.dag.nodes[node_key]
1234 try:
1235 old_state = node[NodeAttributes.STATE]
1236 self._state_map[old_state].remove(node_key)
1237 except KeyError:
1238 if require_old_state:
1239 raise # pragma: no cover
1240 node[NodeAttributes.STATE] = state
1241 node[NodeAttributes.VALUE] = value
1242 self._state_map[state].add(node_key)
1243 self._mark_changed(node_key)
1245 def _set_states(self, node_keys: Iterable[NodeKey], state: States) -> None:
1246 """Set the state of multiple nodes at once.
1248 Materialising the keys is only needed when they have to be read a
1249 second time, which is only when a subscriber will be told about them.
1250 Doing it unconditionally turned the set that callers pass into a tuple,
1251 and ``set.update(tuple)`` rehashes every element where
1252 ``set.update(set)`` reuses the hashes the source set already holds ---
1253 399 extra hashes per insert on a 400-node chain, and the whole of the
1254 subscription API's measured cost to callers who never subscribe.
1255 """
1256 watched = bool(self._subscriptions)
1257 if watched:
1258 node_keys = tuple(node_keys)
1259 for name in node_keys:
1260 node = self.dag.nodes[name]
1261 old_state = node[NodeAttributes.STATE]
1262 self._state_map[old_state].remove(name)
1263 node[NodeAttributes.STATE] = state
1264 self._state_map[state].update(node_keys)
1265 if watched:
1266 # Not via _mark_changed: its star-args would repack the keys again.
1267 self._pending_changed_nodes.update(node_keys)
1269 @_notifies_subscribers()
1270 def set_stale(self, name: Name) -> None:
1271 """Set the state of a node and all its dependencies to STALE.
1273 :param name: Name of the node to set as STALE.
1274 """
1275 node_key = to_nodekey(name)
1276 node_keys: list[NodeKey] = [node_key]
1277 node_keys.extend(nx.dag.descendants(self.dag, node_key))
1278 self._set_states(node_keys, States.STALE)
1279 self._try_set_computable(node_key)
1281 @_notifies_subscribers()
1282 def pin(self, name: Name, value: Any = None) -> None:
1283 """Set the state of a node to PINNED.
1285 :param name: Name of the node to set as PINNED.
1286 :param value: Value to pin to the node, if provided.
1287 :type value: default None
1288 """
1289 node_key = to_nodekey(name)
1290 if value is not None:
1291 self.insert(node_key, value)
1292 self._set_states([node_key], States.PINNED)
1294 @_notifies_subscribers()
1295 def unpin(self, name: Name) -> None:
1296 """Unpin a node (state of node and all descendents will be set to STALE).
1298 :param name: Name of the node to set as PINNED.
1299 """
1300 node_key = to_nodekey(name)
1301 self.set_stale(node_key)
1303 def _get_descendents(self, node_key: NodeKey, stop_states: set[States] | None = None) -> set[NodeKey]:
1304 """Get all descendant nodes, optionally stopping at certain states."""
1305 if stop_states is None:
1306 stop_states = set()
1307 if self.dag.nodes[node_key][NodeAttributes.STATE] in stop_states:
1308 return set()
1309 visited = set()
1310 to_visit = {node_key}
1311 while to_visit:
1312 n = to_visit.pop()
1313 visited.add(n)
1314 for n1 in self.dag.successors(n):
1315 if n1 in visited:
1316 continue
1317 if self.dag.nodes[n1][NodeAttributes.STATE] in stop_states:
1318 continue
1319 to_visit.add(n1)
1320 visited.remove(node_key)
1321 return visited
1323 def _set_descendents(self, node_key: NodeKey, state: States) -> None:
1324 """Set the state of all descendant nodes."""
1325 descendents = self._get_descendents(node_key, {States.PINNED})
1326 self._set_states(descendents, state)
1328 def _set_uninitialized(self, node_key: NodeKey) -> None:
1329 """Set a node to uninitialized state and clear its value."""
1330 self._set_states([node_key], States.UNINITIALIZED)
1331 self.dag.nodes[node_key].pop(NodeAttributes.VALUE, None)
1333 def _set_uptodate(self, node_key: NodeKey, value: object) -> None:
1334 """Set a node to up-to-date state with a value."""
1335 self._set_state_and_value(node_key, States.UPTODATE, value)
1336 self._set_descendents(node_key, States.STALE)
1337 for n in self.dag.successors(node_key):
1338 self._try_set_computable(n)
1340 def _set_error(self, node_key: NodeKey, exc: Exception, tb: str) -> None:
1341 """Set a node to error state with exception information."""
1342 self._set_state_and_literal_value(node_key, States.ERROR, Error(exc, tb))
1343 self._set_descendents(node_key, States.STALE)
1345 def _try_set_computable(self, node_key: NodeKey) -> None:
1346 """Set node to computable if all predecessors are up-to-date."""
1347 if self.dag.nodes[node_key][NodeAttributes.STATE] == States.PINNED:
1348 return
1349 if self.dag.nodes[node_key].get(NodeAttributes.FUNC) is not None:
1350 for n in self.dag.predecessors(node_key):
1351 if not self.dag.has_node(n):
1352 return # pragma: no cover
1353 if self.dag.nodes[n][NodeAttributes.STATE] not in (States.UPTODATE, States.PINNED):
1354 return
1355 self._set_state(node_key, States.COMPUTABLE)
1357 def validate(self) -> ValidationReport:
1358 """Inspect the entire graph for structural and readiness problems.
1360 Validation does not execute functions or mutate the computation.
1361 """
1362 return validate_graph(self.dag, self.executor_map)
1364 def plan(self, targets: Name | Names | None = None) -> ExecutionPlan:
1365 """Describe the work needed to compute one or more targets.
1367 Passing ``None`` plans the whole graph. Planning does not execute
1368 functions or mutate the computation.
1370 :param targets: Target node, list of target nodes, or ``None`` for all nodes.
1371 """
1372 target_node_keys = None if targets is None else names_to_node_keys(targets)
1373 if target_node_keys is not None:
1374 for node_key in target_node_keys:
1375 if not self.dag.has_node(node_key):
1376 msg = f"Node {node_key} does not exist"
1377 raise NonExistentNodeException(msg)
1378 return create_execution_plan(self.dag, self.executor_map, target_node_keys)
1380 def _get_parameter_data(self, node_key: NodeKey) -> Iterable[_ParameterItem]:
1381 """Get all parameter data for a node's function call."""
1382 for arg, value in self.dag.nodes[node_key][NodeAttributes.ARGS].items():
1383 yield _ParameterItem(_ParameterType.ARG, arg, value)
1384 for param_name, value in self.dag.nodes[node_key][NodeAttributes.KWDS].items():
1385 yield _ParameterItem(_ParameterType.KWD, param_name, value)
1386 for in_node_name in self.dag.predecessors(node_key):
1387 param_value = self.dag.nodes[in_node_name][NodeAttributes.VALUE]
1388 edge = self.dag[in_node_name][node_key]
1389 param_type, param_name = edge[EdgeAttributes.PARAM]
1390 yield _ParameterItem(param_type, param_name, param_value)
1392 def _get_func_args_kwds(
1393 self, node_key: NodeKey
1394 ) -> tuple[Callable[..., Any], str | None, list[Any], dict[str, Any]]:
1395 """Get function, executor name, args and kwargs for a node."""
1396 node0 = self.dag.nodes[node_key]
1397 f = node0[NodeAttributes.FUNC]
1398 executor_name = node0.get(NodeAttributes.EXECUTOR)
1399 args: list[Any] = []
1400 kwds: dict[str, Any] = {}
1401 for param in self._get_parameter_data(node_key):
1402 if param.type == _ParameterType.ARG:
1403 idx = param.name
1404 assert isinstance(idx, int) # noqa: S101
1405 while len(args) <= idx:
1406 args.append(None)
1407 args[idx] = param.value
1408 elif param.type == _ParameterType.KWD:
1409 assert isinstance(param.name, str) # noqa: S101
1410 kwds[param.name] = param.value
1411 else: # pragma: no cover
1412 msg = f"Unexpected param type: {param.type}"
1413 raise ValidationError(msg)
1414 return f, executor_name, args, kwds
1416 def get_definition_args_kwds(self, name: Name) -> tuple[list[Any], dict[str, Any]]:
1417 """Get the arguments and keyword arguments for a node's function definition."""
1418 res_args: list[Any] = []
1419 res_kwds: dict[str, Any] = {}
1420 node_key = to_nodekey(name)
1421 node_data = self.dag.nodes[node_key]
1422 if NodeAttributes.ARGS in node_data:
1423 for idx, value in node_data[NodeAttributes.ARGS].items():
1424 while len(res_args) <= idx:
1425 res_args.append(None)
1426 res_args[idx] = C(value)
1427 if NodeAttributes.KWDS in node_data:
1428 for param_name, value in node_data[NodeAttributes.KWDS].items():
1429 res_kwds[param_name] = C(value)
1430 for in_node_key in self.dag.predecessors(node_key):
1431 edge = self.dag[in_node_key][node_key]
1432 if EdgeAttributes.PARAM in edge:
1433 param_type, param_name = edge[EdgeAttributes.PARAM]
1434 if param_type == _ParameterType.ARG:
1435 idx = param_name
1436 assert isinstance(idx, int) # noqa: S101
1437 while len(res_args) <= idx:
1438 res_args.append(None)
1439 res_args[idx] = in_node_key.name
1440 elif param_type == _ParameterType.KWD:
1441 res_kwds[param_name] = in_node_key.name
1442 else: # pragma: no cover
1443 msg = f"Unexpected param type: {param_type}"
1444 raise ValidationError(msg)
1445 return res_args, res_kwds
1447 def _compute_nodes(self, node_keys: Iterable[NodeKey], raise_exceptions: bool = False) -> None:
1448 """Compute multiple nodes, handling dependencies and parallel execution."""
1449 LOG.debug(f"Computing nodes {node_keys}")
1451 futs: dict[Any, NodeKey] = {}
1452 node_keys_set = set(node_keys)
1454 def run(name: NodeKey) -> None:
1455 """Submit a node computation to an executor."""
1456 f, executor_name, args, kwds = self._get_func_args_kwds(name)
1457 executor = self.default_executor if executor_name is None else self.executor_map[executor_name]
1458 fut = executor.submit(_eval_node, name, f, args, kwds, raise_exceptions)
1459 futs[fut] = name
1461 computed: set[NodeKey] = set()
1463 for node_key in node_keys_set:
1464 node0 = self.dag.nodes[node_key]
1465 state = node0[NodeAttributes.STATE]
1466 if state == States.COMPUTABLE:
1467 run(node_key)
1469 while len(futs) > 0:
1470 done, _not_done = wait(futs.keys(), return_when=FIRST_COMPLETED)
1471 for fut in done:
1472 node_key = futs.pop(fut)
1473 node0 = self.dag.nodes[node_key]
1474 try:
1475 value, exc, tb, start_dt, end_dt = fut.result()
1476 except Exception as e:
1477 exc = e
1478 tb = traceback.format_exc()
1479 self._set_error(node_key, exc, tb)
1480 raise
1481 delta = (end_dt - start_dt).total_seconds()
1482 if exc is None:
1483 self._set_state_and_value(node_key, States.UPTODATE, value, throw_conversion_exception=False)
1484 node0[NodeAttributes.TIMING] = TimingData(start_dt, end_dt, delta)
1485 self._set_descendents(node_key, States.STALE)
1486 for n in self.dag.successors(node_key):
1487 logging.debug(str(node_key) + " " + str(n) + " " + str(computed))
1488 if n in computed:
1489 msg = f"Calculating {node_key} for the second time"
1490 raise LoopDetectedException(msg)
1491 self._try_set_computable(n)
1492 node0 = self.dag.nodes[n]
1493 state = node0[NodeAttributes.STATE]
1494 if state == States.COMPUTABLE and n in node_keys_set:
1495 run(n)
1496 else:
1497 assert tb is not None # noqa: S101
1498 self._set_error(node_key, exc, tb)
1499 computed.add(node_key)
1501 def _get_calc_node_keys(self, node_key: NodeKey) -> list[NodeKey]:
1502 """Get node keys that need to be computed for a target node."""
1503 g = nx.DiGraph()
1504 g.add_nodes_from(self.dag.nodes)
1505 g.add_edges_from(self.dag.edges)
1506 for n in nx.ancestors(g, node_key):
1507 node = self.dag.nodes[n]
1508 state = node[NodeAttributes.STATE]
1509 if state == States.UPTODATE or state == States.PINNED:
1510 g.remove_node(n)
1512 ancestors = nx.ancestors(g, node_key)
1513 for n in ancestors:
1514 node = self.dag.nodes[n]
1515 state = node[NodeAttributes.STATE]
1517 if state == States.UNINITIALIZED and len(self.dag.pred[n]) == 0:
1518 msg = f"Cannot compute {node_key} because {n} uninitialized"
1519 raise ValidationError(msg)
1520 if state == States.PLACEHOLDER:
1521 msg = f"Cannot compute {node_key} because {n} is placeholder"
1522 raise ValidationError(msg)
1524 ancestors.add(node_key)
1525 g = g.subgraph(ancestors).copy()
1526 nodes_sorted = topological_sort(g)
1527 return list(nodes_sorted)
1529 def _get_calc_node_names(self, name: Name) -> Names:
1530 """Get node names that need to be computed for a target node."""
1531 node_key = to_nodekey(name)
1532 return node_keys_to_names(self._get_calc_node_keys(node_key))
1534 @_notifies_subscribers()
1535 def compute(self, name: Name | Iterable[Name], raise_exceptions: bool = False) -> None:
1536 """Compute a node or block and all necessary predecessors.
1538 Following the computation, if successful, the target node, and all necessary ancestors that were not already
1539 UPTODATE will have been calculated and set to UPTODATE. Any node that did not need to be calculated will not
1540 have been recalculated.
1542 If any nodes raises an exception, then the state of that node will be set to ERROR, and its value set to an
1543 object containing the exception object, as well as a traceback. This will not halt the computation, which
1544 will proceed as far as it can, until no more nodes that would be required to calculate the target are
1545 COMPUTABLE.
1547 A block name computes every node below that path. Multiple node and block
1548 names may be supplied in a list or generator.
1550 :param name: Name of the node or block to compute
1551 :param raise_exceptions: Whether to pass exceptions raised by node computations back to the caller
1552 :type raise_exceptions: Boolean, default False
1553 """
1554 calc_nodes: set[NodeKey] = set()
1555 names = name if isinstance(name, (types.GeneratorType, list)) else [name]
1556 for name0 in names:
1557 node_key = to_nodekey(name0)
1558 targets = (
1559 [node_key]
1560 if self.has_node(node_key)
1561 else names_to_node_keys(self.get_tree_descendents(node_key, graph_nodes_only=True))
1562 )
1563 if not targets:
1564 targets = [node_key]
1565 for target in targets:
1566 calc_nodes.update(self._get_calc_node_keys(target))
1567 self._compute_nodes(calc_nodes, raise_exceptions=raise_exceptions)
1569 @_notifies_subscribers()
1570 def compute_all(self, raise_exceptions: bool = False) -> None:
1571 """Compute all nodes of a computation that can be computed.
1573 Nodes that are already UPTODATE will not be recalculated. Following the computation, if successful, all
1574 nodes will have state UPTODATE, except UNINITIALIZED input nodes and PLACEHOLDER nodes.
1576 If any nodes raises an exception, then the state of that node will be set to ERROR, and its value set to an
1577 object containing the exception object, as well as a traceback. This will not halt the computation, which
1578 will proceed as far as it can, until no more nodes are COMPUTABLE.
1580 :param raise_exceptions: Whether to pass exceptions raised by node computations back to the caller
1581 :type raise_exceptions: Boolean, default False
1582 """
1583 self._compute_nodes(self._node_keys(), raise_exceptions=raise_exceptions)
1585 def _node_keys(self) -> list[NodeKey]:
1586 """Get a list of nodes in this computation.
1588 :return: List of nodes.
1589 """
1590 return list(self.dag.nodes)
1592 def nodes(self) -> list[Name]:
1593 """Get a list of nodes in this computation.
1595 :return: List of nodes.
1596 """
1597 return [n.name for n in self.dag.nodes]
1599 def get_tree_list_children(self, name: Name) -> set[Name]:
1600 """Get a list of nodes in this computation.
1602 :return: List of nodes.
1603 """
1604 node_key = to_nodekey(name)
1605 idx = len(node_key.parts)
1606 result = set()
1607 for n in self.dag.nodes:
1608 if n.is_descendent_of(node_key):
1609 result.add(n.parts[idx])
1610 return result
1612 def has_node(self, name: Name) -> bool:
1613 """Check if a node with the given name exists in the computation."""
1614 node_key = to_nodekey(name)
1615 return node_key in self.dag.nodes
1617 def tree_has_path(self, name: Name) -> bool:
1618 """Check if a hierarchical path exists in the computation tree."""
1619 node_key = to_nodekey(name)
1620 if node_key.is_root:
1621 return True
1622 if self.has_node(node_key):
1623 return True
1624 return any(n.is_descendent_of(node_key) for n in self.dag.nodes)
1626 def get_tree_descendents(
1627 self, name: Name | None = None, *, include_stem: bool = True, graph_nodes_only: bool = False
1628 ) -> set[Name]:
1629 """Get a list of descendent blocks and nodes.
1631 Returns blocks and nodes that are descendents of the input node,
1632 e.g. for node 'foo', might return ['foo/bar', 'foo/baz'].
1634 :param name: Name of node to get descendents for
1635 :return: List of descendent node names
1636 """
1637 node_key = NodeKey.root() if name is None else to_nodekey(name)
1638 stemsize = len(node_key.parts)
1639 result = set()
1640 for n in self.dag.nodes:
1641 if n.is_descendent_of(node_key):
1642 nodes = [n] if graph_nodes_only else n.ancestors()
1643 for n2 in nodes:
1644 if n2.is_descendent_of(node_key):
1645 nm = n2.name if include_stem else NodeKey(tuple(n2.parts[stemsize:])).name
1646 result.add(nm)
1647 return result
1649 def _state_one(self, name: Name) -> States:
1650 """Get the state of a single node."""
1651 node_key = to_nodekey(name)
1652 state: States = self.dag.nodes[node_key][NodeAttributes.STATE]
1653 return state
1655 @overload
1656 def state(self, name: Name) -> States: ...
1658 @overload
1659 def state(self, name: Names) -> list[States]: ...
1661 def state(self, name: Name | Names) -> States | list[States]:
1662 """Get the state of a node.
1664 This can also be accessed using the attribute-style accessor ``s`` if ``name`` is a valid Python
1665 attribute name::
1667 >>> comp = Computation()
1668 >>> comp.add_node('foo', value=1)
1669 >>> comp.state('foo')
1670 <States.UPTODATE: 4>
1671 >>> comp.s.foo
1672 <States.UPTODATE: 4>
1674 :param name: Name or names of the node to get state for
1675 :type name: Name or Names
1676 """
1677 return apply1(self._state_one, name)
1679 def _value_one(self, name: Name) -> Any:
1680 """Get the value of a single node."""
1681 node_key = to_nodekey(name)
1682 return self.dag.nodes[node_key][NodeAttributes.VALUE]
1684 @overload
1685 def value(self, name: Name) -> Any: ...
1687 @overload
1688 def value(self, name: Names) -> list[Any]: ...
1690 def value(self, name: Name | Names) -> Any | list[Any]:
1691 """Get the current value of a node.
1693 This can also be accessed using the attribute-style accessor ``v`` if ``name`` is a valid Python
1694 attribute name::
1696 >>> comp = Computation()
1697 >>> comp.add_node('foo', value=1)
1698 >>> comp.value('foo')
1699 1
1700 >>> comp.v.foo
1701 1
1703 :param name: Name or names of the node to get the value of
1704 :type name: Name or Names
1705 """
1706 return apply1(self._value_one, name)
1708 def compute_and_get_value(self, name: Name) -> Any:
1709 """Get the current value of a node.
1711 This can also be accessed using the attribute-style accessor ``v`` if ``name`` is a valid Python
1712 attribute name::
1714 >>> comp = Computation()
1715 >>> comp.add_node('foo', value=1)
1716 >>> comp.add_node('bar', lambda foo: foo + 1)
1717 >>> comp.compute_and_get_value('bar')
1718 2
1719 >>> comp.x.bar
1720 2
1722 :param name: Name or names of the node to get the value of
1723 :type name: Name
1724 """
1725 nk = to_nodekey(name)
1726 if self.state(nk) == States.UPTODATE:
1727 return self.value(nk)
1728 self.compute(nk, raise_exceptions=True)
1729 if self.state(nk) == States.UPTODATE:
1730 return self.value(nk)
1731 msg = f"Unable to compute node {nk}"
1732 raise ComputationError(msg)
1734 def _tag_one(self, name: Name) -> set[str]:
1735 """Get the tags of a single node."""
1736 node_key = to_nodekey(name)
1737 node = self.dag.nodes[node_key]
1738 tags: set[str] = node[NodeAttributes.TAG]
1739 return tags
1741 @overload
1742 def tags(self, name: Name) -> set[str]: ...
1744 @overload
1745 def tags(self, name: Names) -> list[set[str]]: ...
1747 def tags(self, name: Name | Names) -> set[str] | list[set[str]]:
1748 """Get the tags associated with a node.
1750 >>> comp = Computation()
1751 >>> comp.add_node('a', tags=['foo', 'bar'])
1752 >>> sorted(comp.t.a)
1753 ['__serialize__', 'bar', 'foo']
1755 :param name: Name or names of the node to get the tags of
1756 :return:
1757 """
1758 return apply1(self._tag_one, name)
1760 def nodes_by_tag(self, tag: str | Iterable[str]) -> set[Name]:
1761 """Get the names of nodes with a particular tag or tags.
1763 :param tag: Tag or tags for which to retrieve nodes
1764 :return: Names of the nodes with those tags
1765 """
1766 nodes: set[NodeKey] = set()
1767 tags_to_check: Iterable[str] = [tag] if isinstance(tag, str) else tag
1768 for tag1 in tags_to_check:
1769 nodes1 = self._tag_map.get(tag1)
1770 if nodes1 is not None:
1771 nodes.update(nodes1)
1772 return {n.name for n in nodes}
1774 def _style_one(self, name: Name) -> str | None:
1775 """Get the style of a single node."""
1776 node_key = to_nodekey(name)
1777 node = self.dag.nodes[node_key]
1778 style: str | None = node.get(NodeAttributes.STYLE)
1779 return style
1781 @overload
1782 def styles(self, name: Name) -> str | None: ...
1784 @overload
1785 def styles(self, name: Names) -> list[str | None]: ...
1787 def styles(self, name: Name | Names) -> str | list[str | None] | None:
1788 """Get the tags associated with a node.
1790 >>> comp = Computation()
1791 >>> comp.add_node('a', style='dot')
1792 >>> comp.style.a
1793 'dot'
1795 :param name: Name or names of the node to get the tags of
1796 :return:
1797 """
1798 return apply1(self._style_one, name)
1800 def _get_item_one(self, name: Name) -> NodeData:
1801 """Get state and value data for a single node."""
1802 node_key = to_nodekey(name)
1803 node = self.dag.nodes[node_key]
1804 return NodeData(node[NodeAttributes.STATE], node[NodeAttributes.VALUE])
1806 @overload
1807 def __getitem__(self, name: Name) -> NodeData: ...
1809 @overload
1810 def __getitem__(self, name: Names) -> list[NodeData]: ...
1812 def __getitem__(self, name: Name | Names) -> NodeData | list[NodeData]:
1813 """Get the state and current value of a node.
1815 :param name: Name of the node to get the state and value of
1816 """
1817 return apply1(self._get_item_one, name)
1819 def _get_timing_one(self, name: Name) -> TimingData | None:
1820 """Get timing data for a single node."""
1821 node_key = to_nodekey(name)
1822 node = self.dag.nodes[node_key]
1823 timing: TimingData | None = node.get(NodeAttributes.TIMING, None)
1824 return timing
1826 @overload
1827 def get_timing(self, name: Name) -> TimingData | None: ...
1829 @overload
1830 def get_timing(self, name: Names) -> list[TimingData | None]: ...
1832 def get_timing(self, name: Name | Names) -> TimingData | list[TimingData | None] | None:
1833 """Get the timing information for a node.
1835 :param name: Name or names of the node to get the timing information of
1836 :return:
1837 """
1838 return apply1(self._get_timing_one, name)
1840 def to_df(self) -> pd.DataFrame:
1841 """Get a dataframe containing the states and value of all nodes of computation.
1843 ::
1845 >>> import loman
1846 >>> comp = loman.Computation()
1847 >>> comp.add_node('foo', value=1)
1848 >>> comp.add_node('bar', value=2)
1849 >>> comp.to_df() # doctest: +NORMALIZE_WHITESPACE
1850 state value
1851 foo States.UPTODATE 1
1852 bar States.UPTODATE 2
1853 """
1854 df = pd.DataFrame(index=topological_sort(self.dag))
1855 df[NodeAttributes.STATE] = pd.Series(nx.get_node_attributes(self.dag, NodeAttributes.STATE))
1856 df[NodeAttributes.VALUE] = pd.Series(nx.get_node_attributes(self.dag, NodeAttributes.VALUE))
1857 df_timing = pd.DataFrame.from_dict(nx.get_node_attributes(self.dag, "timing"), orient="index")
1858 df = pd.merge(df, df_timing, left_index=True, right_index=True, how="left")
1859 df.index = pd.Index([nk.name for nk in df.index])
1860 return df
1862 def to_dict(self) -> dict[NodeKey, Any]:
1863 """Get a dictionary containing the values of all nodes of a computation.
1865 ::
1867 >>> import loman
1868 >>> comp = loman.Computation()
1869 >>> comp.add_node('foo', value=1)
1870 >>> comp.add_node('bar', value=2)
1871 >>> comp.to_dict() # doctest: +ELLIPSIS
1872 {NodeKey('foo'): 1, NodeKey('bar'): 2}
1873 """
1874 result: dict[NodeKey, Any] = nx.get_node_attributes(self.dag, NodeAttributes.VALUE)
1875 return result
1877 def _get_inputs_one_node_keys(self, node_key: NodeKey) -> list[NodeKey | None]:
1878 """Get input node keys for a single node."""
1879 args_dict: dict[int, NodeKey] = {}
1880 kwds: list[NodeKey | None] = []
1881 max_arg_index = -1
1882 for input_node in self.dag.predecessors(node_key):
1883 input_edge = self.dag[input_node][node_key]
1884 input_type, input_param = input_edge[EdgeAttributes.PARAM]
1885 if input_type == _ParameterType.ARG:
1886 idx = input_param
1887 max_arg_index = max(max_arg_index, idx)
1888 args_dict[idx] = input_node
1889 elif input_type == _ParameterType.KWD:
1890 kwds.append(input_node)
1891 if max_arg_index >= 0:
1892 args: list[NodeKey | None] = [None] * (max_arg_index + 1)
1893 for idx, input_node in args_dict.items():
1894 args[idx] = input_node
1895 result: list[NodeKey | None] = args + kwds
1896 return result
1897 else:
1898 return kwds
1900 def _get_inputs_one_names(self, name: Name) -> Names:
1901 """Get input node names for a single node."""
1902 node_key = to_nodekey(name)
1903 return node_keys_to_names([nk for nk in self._get_inputs_one_node_keys(node_key) if nk is not None])
1905 @overload
1906 def get_inputs(self, name: Name) -> Names: ...
1908 @overload
1909 def get_inputs(self, name: Names) -> list[Names]: ...
1911 def get_inputs(self, name: Name | Names) -> Names | list[Names]:
1912 """Get a list of the inputs for a node or set of nodes.
1914 :param name: Name or names of nodes to get inputs for
1915 :return: If name is scalar, return a list of upstream nodes used as input. If name is a list, return a
1916 list of list of inputs.
1917 """
1918 return apply1(self._get_inputs_one_names, name)
1920 def _get_ancestors_node_keys(self, node_keys: Iterable[NodeKey], include_self: bool = True) -> set[NodeKey]:
1921 """Get all ancestor node keys for a set of nodes."""
1922 ancestors: set[NodeKey] = set()
1923 for n in node_keys:
1924 if include_self:
1925 ancestors.add(n)
1926 for ancestor in nx.ancestors(self.dag, n):
1927 ancestors.add(ancestor)
1928 return ancestors
1930 def get_ancestors(self, names: Name | Names, include_self: bool = True) -> Names:
1931 """Get all ancestor nodes of the specified nodes."""
1932 node_keys = names_to_node_keys(names)
1933 ancestor_node_keys = self._get_ancestors_node_keys(node_keys, include_self)
1934 return node_keys_to_names(ancestor_node_keys)
1936 def _get_original_inputs_node_keys(self, node_keys: list[NodeKey] | None) -> list[NodeKey]:
1937 """Get original input node keys that have no computation function."""
1938 resolved_node_keys: Iterable[NodeKey]
1939 resolved_node_keys = self._node_keys() if node_keys is None else self._get_ancestors_node_keys(node_keys)
1940 return [n for n in resolved_node_keys if self.dag.nodes[n].get(NodeAttributes.FUNC) is None]
1942 def get_original_inputs(self, names: Name | Names | None = None) -> Names:
1943 """Get a list of the original non-computed inputs for a node or set of nodes.
1945 :param names: Name or names of nodes to get inputs for
1946 :return: Return a list of original non-computed inputs that are ancestors of the input nodes
1947 """
1948 nks = None if names is None else names_to_node_keys(names)
1950 result_nks = self._get_original_inputs_node_keys(nks)
1952 return node_keys_to_names(result_nks)
1954 def _get_outputs_one(self, name: Name) -> Names:
1955 """Get output node names for a single node."""
1956 node_key = to_nodekey(name)
1957 output_node_keys = list(self.dag.successors(node_key))
1958 return node_keys_to_names(output_node_keys)
1960 @overload
1961 def get_outputs(self, name: Name) -> Names: ...
1963 @overload
1964 def get_outputs(self, name: Names) -> list[Names]: ...
1966 def get_outputs(self, name: Name | Names) -> Names | list[Names]:
1967 """Get a list of the outputs for a node or set of nodes.
1969 :param name: Name or names of nodes to get outputs for
1970 :return: If name is scalar, return a list of downstream nodes used as output. If name is a list, return a
1971 list of list of outputs.
1973 """
1974 return apply1(self._get_outputs_one, name)
1976 def _get_descendents_node_keys(self, node_keys: Iterable[NodeKey], include_self: bool = True) -> set[NodeKey]:
1977 """Get all descendant node keys for a set of nodes."""
1978 descendent_node_keys: set[NodeKey] = set()
1979 for node_key in node_keys:
1980 if include_self:
1981 descendent_node_keys.add(node_key)
1982 for descendent in nx.descendants(self.dag, node_key):
1983 descendent_node_keys.add(descendent)
1984 return descendent_node_keys
1986 def get_descendents(self, names: Name | Names, include_self: bool = True) -> Names:
1987 """Get all descendent nodes of the specified nodes."""
1988 node_keys = names_to_node_keys(names)
1989 descendent_node_keys = self._get_descendents_node_keys(node_keys, include_self)
1990 return node_keys_to_names(descendent_node_keys)
1992 def get_final_outputs(self, names: Name | Names | None = None) -> Names:
1993 """Get final output nodes (nodes with no descendants) from the specified nodes."""
1994 final_node_keys: Iterable[NodeKey]
1995 if names is None:
1996 final_node_keys = self._node_keys()
1997 else:
1998 nks = names_to_node_keys(names)
1999 final_node_keys = self._get_descendents_node_keys(nks)
2000 output_node_keys = [n for n in final_node_keys if len(nx.descendants(self.dag, n)) == 0]
2001 return node_keys_to_names(output_node_keys)
2003 def get_source(self, name: Name) -> str:
2004 """Get the source code for a node."""
2005 node_key = to_nodekey(name)
2006 func = self.dag.nodes[node_key].get(NodeAttributes.FUNC, None)
2007 if func is not None:
2008 file = inspect.getsourcefile(func)
2009 _, lineno = inspect.getsourcelines(func)
2010 source = inspect.getsource(func)
2011 return f"{file}:{lineno}\n\n{source}"
2012 else:
2013 return "NOT A CALCULATED NODE"
2015 def print_source(self, name: Name) -> None:
2016 """Print the source code for a computation node."""
2017 print(self.get_source(name))
2019 @_notifies_subscribers(graph_changed=True)
2020 def restrict(self, output_names: Name | Names, input_names: Name | Names | None = None) -> None:
2021 """Restrict a computation to the ancestors of a set of output nodes.
2023 Excludes ancestors of a set of input nodes.
2025 If the set of input_nodes that is specified is not sufficient for the set of output_nodes then additional
2026 nodes that are ancestors of the output_nodes will be included, but the input nodes specified will be input
2027 nodes of the modified Computation.
2029 :param output_nodes:
2030 :param input_nodes:
2031 :return: None - modifies existing computation in place
2032 """
2033 if input_names is not None:
2034 for n in as_iterable(input_names):
2035 nodedata = self._get_item_one(n)
2036 self.add_node(n)
2037 self._set_state_and_literal_value(to_nodekey(n), nodedata.state, nodedata.value)
2038 output_node_keys = names_to_node_keys(output_names)
2039 ancestor_node_keys = self._get_ancestors_node_keys(output_node_keys)
2040 removed = [n for n in self.dag if n not in ancestor_node_keys]
2041 self.dag.remove_nodes_from(removed)
2042 self._mark_changed(*removed)
2044 def __getstate__(self) -> dict[str, Any]:
2045 """Prepare computation for serialization by removing non-serializable nodes."""
2046 node_serialize = nx.get_node_attributes(self.dag, NodeAttributes.TAG)
2047 obj = self.copy()
2048 for name, tags in node_serialize.items():
2049 if SystemTags.SERIALIZE not in tags:
2050 obj._set_uninitialized(name)
2051 return {"dag": obj.dag}
2053 def __setstate__(self, state: dict[str, Any]) -> None:
2054 """Restore computation from serialized state."""
2055 self.__init__()
2056 self.dag = state["dag"]
2057 self._refresh_maps()
2059 def write_dill_old(self, file_: str | BinaryIO) -> None:
2060 """Serialize a computation to a file or file-like object.
2062 .. deprecated::
2063 Superseded by :meth:`write_dill`, and by :meth:`save` in turn. Kept
2064 because removing it would break callers without warning; it will go
2065 in a release that says so.
2067 .. warning::
2068 Not safe to call concurrently. It removes ``__getstate__`` and
2069 ``__setstate__`` from the class for the duration of the write, which
2070 is process-wide, so another thread pickling a Computation at the same
2071 moment gets the wrong representation. :meth:`write_dill` has no such
2072 problem, and :meth:`save` supersedes both.
2074 :param file_: If string, writes to a file
2075 :type file_: File-like object, or string
2076 """
2077 warnings.warn(
2078 "write_dill_old is deprecated and will be removed in a future release. Use save instead.",
2079 DeprecationWarning,
2080 stacklevel=2,
2081 )
2082 original_getstate = self.__class__.__getstate__
2083 original_setstate = self.__class__.__setstate__
2085 try:
2086 del self.__class__.__getstate__
2087 del self.__class__.__setstate__
2089 node_serialize = nx.get_node_attributes(self.dag, NodeAttributes.TAG)
2090 obj = self.copy()
2091 obj.executor_map = None # type: ignore[assignment]
2092 obj.default_executor = None # type: ignore[assignment]
2093 for name, tags in node_serialize.items():
2094 if SystemTags.SERIALIZE not in tags:
2095 obj._set_uninitialized(name)
2097 if isinstance(file_, str):
2098 with open(file_, "wb") as f:
2099 dill.dump(obj, f)
2100 else:
2101 dill.dump(obj, file_)
2102 finally:
2103 self.__class__.__getstate__ = original_getstate # type: ignore[method-assign]
2104 self.__class__.__setstate__ = original_setstate
2106 def write_dill(self, file_: str | BinaryIO) -> None:
2107 """Serialize a computation to a file or file-like object.
2109 .. deprecated::
2110 Use :meth:`write_json` instead. dill-based serialization will be
2111 removed in a future release.
2113 :param file_: If string, writes to a file
2114 :type file_: File-like object, or string
2115 """
2116 warnings.warn(
2117 "write_dill is deprecated and will be removed in a future release. Use write_json instead.",
2118 DeprecationWarning,
2119 stacklevel=2,
2120 )
2121 if isinstance(file_, str):
2122 with open(file_, "wb") as f:
2123 dill.dump(self, f)
2124 else:
2125 dill.dump(self, file_)
2127 @staticmethod
2128 def read_dill(file_: str | BinaryIO) -> "Computation":
2129 """Deserialize a computation from a file or file-like object.
2131 .. deprecated::
2132 Use :meth:`read_json` instead. dill-based serialization will be
2133 removed in a future release.
2135 .. warning::
2136 This method uses dill.load() which can execute arbitrary code.
2137 Only load files from trusted sources. Never load data from
2138 untrusted or unauthenticated sources as it may lead to arbitrary
2139 code execution.
2141 :param file_: If string, writes to a file
2142 :type file_: File-like object, or string
2143 """
2144 warnings.warn(
2145 "read_dill is deprecated and will be removed in a future release. Use read_json instead.",
2146 DeprecationWarning,
2147 stacklevel=2,
2148 )
2149 if isinstance(file_, str):
2150 with open(file_, "rb") as f:
2151 obj = dill.load(f) # noqa: S301 # nosec B301
2152 else:
2153 obj = dill.load(file_) # noqa: S301 # nosec B301
2154 if isinstance(obj, Computation):
2155 return obj
2156 else:
2157 msg = "Loaded object is not a Computation"
2158 raise ValidationError(msg)
2160 def save(
2161 self,
2162 path: str,
2163 *,
2164 profile: "str | SerializationProfile | None" = None,
2165 container: str | None = None,
2166 stores: "dict[str, BlobStore] | None" = None,
2167 serializer: "ComputationSerializer | None" = None,
2168 ) -> None:
2169 """Save this computation to *path*.
2171 The default is a ``.loman`` file: one zip holding a ``manifest.json``
2172 describing the graph, plus a ``blobs/`` directory holding large values as
2173 binary. The manifest still records every value's shape, dtype and index
2174 type, so the file can be inspected without decoding any of the data.
2176 ::
2178 comp.save('run.loman') # efficient, zipped
2179 comp.save('run.loman', profile='readable') # inline JSON, zipped
2180 comp.save('run.json') # single JSON document
2181 comp.save('run_dir', container='dir') # same layout, unzipped
2183 *profile* and *container* are independent. The profile decides whether a
2184 value's bytes are written inline or out of line; the container decides
2185 where they land. The one combination that cannot work is the efficient
2186 profile in the ``json`` container, which raises.
2188 Prefer ``container='dir'`` when saving repeatedly --- updating one value
2189 in a zip rewrites the whole archive, at a cost that grows with its size,
2190 while a directory rewrites only the file that changed.
2192 :param path: Destination path. A ``.json`` suffix selects the single
2193 document container; anything else defaults to a ``.loman`` zip.
2194 :param profile: ``"readable"``, ``"efficient"`` (the default), or a
2195 :class:`~loman.serialization.profile.SerializationProfile`.
2196 :param container: ``"zip"``, ``"dir"`` or ``"json"``. Inferred from
2197 *path* when omitted.
2198 :param stores: Named
2199 :class:`~loman.serialization.blobs.BlobStore` instances for values
2200 that belong somewhere other than the saved file --- a bucket, a
2201 database. A node is routed to one by ``add_node(store=...)`` or a
2202 profile override. The same names must be supplied to :meth:`load`.
2203 :param serializer: Optional custom serializer, for user-defined types.
2204 """
2205 from .serialization.computation import ComputationSerializer
2207 s = serializer if serializer is not None else ComputationSerializer()
2208 s.save(self, path, profile=profile, container=container, stores=stores)
2210 @staticmethod
2211 def load(
2212 path: str,
2213 *,
2214 serializer: "ComputationSerializer | None" = None,
2215 allow_code: bool = True,
2216 stores: "dict[str, BlobStore] | None" = None,
2217 ) -> "Computation":
2218 """Load a computation saved by :meth:`save`, in any container.
2220 The container is detected from the file itself, so a ``.loman`` archive,
2221 a directory and a plain JSON document all load through this one call ---
2222 including documents written by earlier format versions.
2224 .. warning::
2225 Loading restores node functions, which means importing the modules
2226 the file names, or unpickling a dill blob out of it. Both run code
2227 chosen by the file. Only load files you trust, or pass
2228 ``allow_code=False``.
2230 :param path: Path to a ``.loman`` file, a container directory, or a
2231 JSON document.
2232 :param serializer: Optional custom serializer, matching the one used to
2233 save.
2234 :param allow_code: When false, callables are not resolved; values,
2235 structure, states and tags still load.
2236 :param stores: Named stores for values held outside the file. A saved
2237 file records a store's name but never its configuration, so it never
2238 contains a bucket or a credential --- and cannot resolve external
2239 values without the matching store being supplied here.
2240 :rtype: Computation
2241 """
2242 from .serialization.computation import ComputationSerializer
2244 return ComputationSerializer.load_path(path, serializer=serializer, allow_code=allow_code, stores=stores)
2246 def write_json(self, file_: str | TextIO, *, serializer: "ComputationSerializer | None" = None) -> None:
2247 """Serialize a computation to a JSON file or file-like object.
2249 Custom types can be supported by passing a custom *serializer* —
2250 either a :class:`~loman.serialization.computation.ComputationSerializer`
2251 instance with extra transformers registered, or a subclass that
2252 overrides the transformer factory.
2254 :param file_: Destination file path (str) or text-mode file-like object.
2255 :param serializer: Optional custom serializer. If ``None`` the default
2256 :class:`~loman.serialization.computation.ComputationSerializer` is used.
2257 """
2258 from .serialization.computation import ComputationSerializer
2260 s = serializer if serializer is not None else ComputationSerializer()
2261 if isinstance(file_, str):
2262 with open(file_, "w", encoding="utf-8") as f:
2263 s.dump(self, f)
2264 else:
2265 s.dump(self, file_)
2267 @staticmethod
2268 def read_json(
2269 file_: str | TextIO,
2270 *,
2271 serializer: "ComputationSerializer | None" = None,
2272 allow_code: bool = True,
2273 ) -> "Computation":
2274 """Deserialize a computation from a JSON file or file-like object.
2276 .. warning::
2277 Loading a computation restores its node functions, which means
2278 importing the modules the file names, or unpickling a dill blob out
2279 of it. Both run code chosen by the file. Only load files from
2280 sources you trust, or pass ``allow_code=False``.
2282 :param file_: Source file path (str) or text-mode file-like object.
2283 :param serializer: Optional custom serializer. If ``None`` the default
2284 :class:`~loman.serialization.computation.ComputationSerializer` is used.
2285 :param allow_code: When false, encoded callables are skipped rather than
2286 resolved, so no module named by the file is imported and no dill blob
2287 is unpickled. Values, structure, states and tags still load; every
2288 node's function and converter comes back as ``None``, so the graph
2289 can be inspected but not recalculated.
2290 :rtype: Computation
2291 """
2292 from .serialization.computation import ComputationSerializer
2294 s = serializer if serializer is not None else ComputationSerializer()
2295 if isinstance(file_, str):
2296 with open(file_, encoding="utf-8") as f:
2297 return s.load(f, allow_code=allow_code)
2298 else:
2299 return s.load(file_, allow_code=allow_code)
2301 def copy(self) -> "Computation":
2302 """Create a copy of a computation.
2304 The copy is shallow. Any values in the new Computation's DAG will be the same object as this Computation's
2305 DAG. As new objects will be created by any further computations, this should not be an issue.
2307 :rtype: Computation
2308 """
2309 obj = Computation()
2310 obj.dag = nx.DiGraph(self.dag)
2311 obj._tag_map = defaultdict(set, {tag: nodes.copy() for tag, nodes in self._tag_map.items()})
2312 obj._state_map = {state: nodes.copy() for state, nodes in self._state_map.items()}
2313 return obj
2315 @_notifies_subscribers(graph_changed=True)
2316 def add_named_tuple_expansion(self, name: Name, namedtuple_type: type, group: str | None = None) -> None:
2317 """Automatically add nodes to extract each element of a named tuple type.
2319 It is often convenient for a calculation to return multiple values, and it is polite to do this a namedtuple
2320 rather than a regular tuple, so that later users have same name to identify elements of the tuple. It can
2321 also help make a computation clearer if a downstream computation depends on one element of such a tuple,
2322 rather than the entire tuple. This does not affect the computation per se, but it does make the intention
2323 clearer.
2325 To avoid having to create many boiler-plate node definitions to expand namedtuples, the
2326 ``add_named_tuple_expansion`` method automatically creates new nodes for each element of a tuple. The
2327 convention is that an element called 'element', in a node called 'node' will be expanded into a new node
2328 called 'node.element', and that this will be applied for each element.
2330 Example::
2332 >>> from collections import namedtuple
2333 >>> Coordinate = namedtuple('Coordinate', ['x', 'y'])
2334 >>> comp = Computation()
2335 >>> comp.add_node('c', value=Coordinate(1, 2))
2336 >>> comp.add_named_tuple_expansion('c', Coordinate)
2337 >>> comp.compute_all()
2338 >>> comp.value('c.x')
2339 1
2340 >>> comp.value('c.y')
2341 2
2343 :param name: Node to cera
2344 :param namedtuple_type: Expected type of the node
2345 :type namedtuple_type: namedtuple class
2346 """
2348 def make_f(field_name: str) -> Callable[[Any], Any]:
2349 """Create a function to extract a field from a namedtuple."""
2351 def get_field_value(tuple_val: Any) -> Any:
2352 """Extract field value from the namedtuple."""
2353 return getattr(tuple_val, field_name)
2355 return get_field_value
2357 for field_name in namedtuple_type._fields: # type: ignore[attr-defined]
2358 node_name = f"{name}.{field_name}"
2359 self.add_node(node_name, make_f(field_name), kwds={"tuple_val": name}, group=group)
2360 self.set_tag(node_name, SystemTags.EXPANSION)
2362 @_notifies_subscribers(graph_changed=True)
2363 def add_map_node(
2364 self,
2365 result_node: Name,
2366 input_node: Name,
2367 subgraph: "Computation",
2368 subgraph_input_node: Name,
2369 subgraph_output_node: Name,
2370 ) -> None:
2371 """Apply a graph to each element of iterable.
2373 In turn, each element in the ``input_node`` of this graph will be inserted in turn into the subgraph's
2374 ``subgraph_input_node``, then the subgraph's ``subgraph_output_node`` calculated. The resultant list, with
2375 an element or each element in ``input_node``, will be inserted into ``result_node`` of this graph. In this
2376 way ``add_map_node`` is similar to ``map`` in functional programming.
2378 :param result_node: The node to place a list of results in **this** graph
2379 :param input_node: The node to get a list input values from **this** graph
2380 :param subgraph: The graph to use to perform calculation for each element
2381 :param subgraph_input_node: The node in **subgraph** to insert each element in turn
2382 :param subgraph_output_node: The node in **subgraph** to read the result for each element
2383 """
2385 def f(xs: Iterable[Any]) -> list[Any]:
2386 """Apply subgraph computation to each element in the input."""
2387 results: list[Any] = []
2388 is_error = False
2389 for x in xs:
2390 subgraph.insert(subgraph_input_node, x)
2391 subgraph.compute(subgraph_output_node)
2392 if subgraph.state(subgraph_output_node) == States.UPTODATE:
2393 results.append(subgraph.value(subgraph_output_node))
2394 else:
2395 is_error = True
2396 results.append(subgraph.copy())
2397 if is_error:
2398 msg = f"Unable to calculate {result_node}"
2399 raise MapException(msg, results)
2400 return results
2402 self.add_node(result_node, f, kwds={"xs": input_node})
2404 def prepend_path(self, path: Name | ConstantValue, prefix_path: NodeKey) -> NodeKey | ConstantValue:
2405 """Prepend a prefix path to a node path."""
2406 if isinstance(path, ConstantValue):
2407 return path
2408 nk = to_nodekey(path)
2409 return prefix_path.join(nk)
2411 @_notifies_subscribers(graph_changed=True)
2412 def add_block(
2413 self,
2414 base_path: Name,
2415 block: "Computation",
2416 *,
2417 keep_values: bool = True,
2418 links: dict[str, Name] | None = None,
2419 metadata: dict[str, Any] | None = None,
2420 ) -> None:
2421 """Add a computation block as a subgraph to this computation.
2423 ``keep_values`` defaults to ``True``, so values already held by ``block``
2424 are copied along with its structure: a block added this way is often a
2425 sub-model that has already been populated or calibrated, and would not be
2426 computable without them. The repeated-block utilities in
2427 :mod:`loman.util` default to ``False`` instead, because they stamp out
2428 many copies of one template. See :class:`loman.util.RepeatedBlocks`.
2430 :param base_path: Parent path to add the block's nodes below
2431 :param block: Computation to copy into this computation
2432 :param keep_values: Whether to copy the block's current values
2433 :param links: Mapping from the block's relative input names to outer nodes
2434 :param metadata: Metadata to attach to the block path
2435 """
2436 base_path_nk = to_nodekey(base_path)
2437 for node_name in block.nodes():
2438 node_key = to_nodekey(node_name)
2439 node_data = block.dag.nodes[node_key]
2440 tags = node_data.get(NodeAttributes.TAG, None)
2441 # strip the serialize tag from the original node: add_block explicitly
2442 # sets serialize=False, meaning "don't serialize the function".
2443 # Value serialization is controlled separately via keep_values below.
2444 if tags is not None:
2445 tags = tags - {SystemTags.SERIALIZE}
2446 style = node_data.get(NodeAttributes.STYLE, None)
2447 group = node_data.get(NodeAttributes.GROUP, None)
2448 args_def, kwds_def = block.get_definition_args_kwds(node_key)
2449 args_prepended = [self.prepend_path(arg, base_path_nk) for arg in args_def]
2450 kwds_prepended = {k: self.prepend_path(v, base_path_nk) for k, v in kwds_def.items()}
2451 func = node_data.get(NodeAttributes.FUNC, None)
2452 executor = node_data.get(NodeAttributes.EXECUTOR, None)
2453 converter = node_data.get(NodeAttributes.CONVERTER, None)
2454 new_node_name = self.prepend_path(node_name, base_path_nk)
2455 self.add_node(
2456 new_node_name,
2457 func,
2458 args=args_prepended,
2459 kwds=kwds_prepended,
2460 converter=converter,
2461 serialize=False,
2462 inspect=False,
2463 group=group,
2464 tags=tags,
2465 style=style,
2466 executor=executor,
2467 )
2468 if keep_values and NodeAttributes.VALUE in node_data:
2469 new_node_key = to_nodekey(new_node_name)
2470 self._set_state_and_literal_value(
2471 new_node_key, node_data[NodeAttributes.STATE], node_data[NodeAttributes.VALUE]
2472 )
2473 # The node has a concrete value — mark it serializable so the
2474 # value survives a JSON roundtrip even though the function is not.
2475 self._set_tag_one(new_node_key, SystemTags.SERIALIZE)
2476 if links is not None:
2477 for target, source in links.items():
2478 self.link(base_path_nk.join_parts(target), source)
2479 if metadata is not None:
2480 self._metadata[base_path_nk] = metadata
2481 else:
2482 if base_path_nk in self._metadata:
2483 del self._metadata[base_path_nk]
2485 @_notifies_subscribers(graph_changed=True)
2486 def link(self, target: Name, source: Name) -> None:
2487 """Create a link between two nodes in the computation graph."""
2488 target_nk = to_nodekey(target)
2489 source_nk = to_nodekey(source)
2490 if target_nk == source_nk:
2491 return
2493 target_style = self._style_one(target_nk) if self.has_node(target_nk) else None
2494 source_style = self._style_one(source_nk) if self.has_node(source_nk) else None
2495 style = target_style if target_style else source_style
2497 self.add_node(target_nk, identity_function, kwds={"x": source_nk}, style=style)
2499 def _repr_svg_(self) -> str | None:
2500 """Return SVG representation for Jupyter notebook display."""
2501 return GraphView(self).svg()
2503 def draw(
2504 self,
2505 root: NodeKey | None = None,
2506 *,
2507 node_transformations: dict[Name, str] | None = None,
2508 cmap: Any = None,
2509 colors: str = "state",
2510 shapes: str | None = None,
2511 graph_attr: dict[str, Any] | None = None,
2512 node_attr: dict[str, Any] | None = None,
2513 edge_attr: dict[str, Any] | None = None,
2514 show_expansion: bool = False,
2515 collapse_all: bool = True,
2516 ) -> GraphView:
2517 """Draw a computation's current state using the GraphViz utility.
2519 :param root: Optional PathType. Sub-block to draw
2520 :param cmap: Default: None
2521 :param colors: 'state' - colors indicate state. 'timing' - colors indicate execution time. Default: 'state'.
2522 :param shapes: None - ovals. 'type' - shapes indicate type. Default: None.
2523 :param graph_attr: Mapping of (attribute, value) pairs for the graph. For example
2524 ``graph_attr={'size': '"10,8"'}`` can control the size of the output graph
2525 :param node_attr: Mapping of (attribute, value) pairs set for all nodes.
2526 :param edge_attr: Mapping of (attribute, value) pairs set for all edges.
2527 :param collapse_all: Whether to collapse all blocks that aren't explicitly expanded.
2528 """
2529 node_formatter = NodeFormatter.create(cmap, colors, shapes)
2530 node_transformations_copy: dict[Name, str] = (
2531 node_transformations.copy() if node_transformations is not None else {}
2532 )
2533 if not show_expansion:
2534 for nodekey in self.nodes_by_tag(SystemTags.EXPANSION):
2535 node_transformations_copy[nodekey] = NodeTransformations.CONTRACT
2536 v = GraphView(
2537 self,
2538 root=root,
2539 node_formatter=node_formatter,
2540 graph_attr=graph_attr,
2541 node_attr=node_attr,
2542 edge_attr=edge_attr,
2543 node_transformations=node_transformations_copy,
2544 collapse_all=collapse_all,
2545 )
2546 return v
2548 def widget(
2549 self,
2550 root: NodeKey | None = None,
2551 *,
2552 node_transformations: dict[Name, str] | None = None,
2553 cmap: Any = None,
2554 colors: str = "state",
2555 shapes: str | None = None,
2556 graph_attr: dict[str, Any] | None = None,
2557 node_attr: dict[str, Any] | None = None,
2558 edge_attr: dict[str, Any] | None = None,
2559 show_expansion: bool = False,
2560 collapse_all: bool = True,
2561 editable: bool = True,
2562 buildable: bool = False,
2563 namespace: dict[str, Any] | None = None,
2564 fit_on_render: bool = False,
2565 max_rendered_nodes: int = 500,
2566 rankdir: str = "LR",
2567 ) -> "ComputationWidget":
2568 """Create an interactive notebook widget for this computation.
2570 Requires the ``ui`` extra: ``pip install 'loman[ui]'``. It is imported
2571 lazily, so ordinary use of Loman does not load AnyWidget or its notebook
2572 dependencies.
2574 Arguments other than ``editable`` mirror :meth:`draw`, so the two are
2575 learnable together. The widget subscribes to this computation and
2576 follows it; ``comp.draw()`` and ``_repr_svg_`` stay static pictures.
2578 Note that ``colors="state"`` repaints in place, while any other
2579 colouring re-runs Graphviz on every change. See
2580 :class:`~loman.ui.ComputationWidget` for the details.
2582 :param editable: Allow scalar input editing and computation controls.
2583 Expanding and collapsing blocks stays available either way.
2584 :param buildable: Allow the graph itself to be built in the
2585 widget: adding, redefining, renaming and deleting nodes. Off by
2586 default, and needs ``editable`` too, because defining a calculation
2587 node compiles and runs an expression typed in the browser inside
2588 this kernel.
2589 :param namespace: Globals a node built in the widget is compiled
2590 against, so its expression can use the notebook's own imports.
2591 Pass ``globals()``. The default is an empty namespace, in which
2592 only builtins are in scope.
2593 :param fit_on_render: Scale the graph to fit the pane on every render,
2594 rather than opening at natural size. Useful when the shape of a
2595 large graph matters more than its labels.
2596 :param max_rendered_nodes: Refuse to open a block that would put more
2597 than this many nodes on screen. Does not cap the initial view.
2598 :param rankdir: Initial Graphviz layout direction, ``LR`` (default) or
2599 ``TB``, toggled live from the toolbar.
2600 :return: A live widget subscribed to this computation.
2601 """
2602 from .ui import ComputationWidget
2604 return ComputationWidget(
2605 self,
2606 root=root,
2607 node_transformations=node_transformations,
2608 cmap=cmap,
2609 colors=colors,
2610 shapes=shapes,
2611 graph_attr=graph_attr,
2612 node_attr=node_attr,
2613 edge_attr=edge_attr,
2614 show_expansion=show_expansion,
2615 collapse_all=collapse_all,
2616 editable=editable,
2617 buildable=buildable,
2618 namespace=namespace,
2619 fit_on_render=fit_on_render,
2620 max_rendered_nodes=max_rendered_nodes,
2621 rankdir=rankdir,
2622 )
2624 def view(self, cmap: Any = None, colors: str = "state", shapes: str | None = None) -> None:
2625 """Create and display a visualization of the computation graph."""
2626 node_formatter = NodeFormatter.create(cmap, colors, shapes)
2627 v = GraphView(self, node_formatter=node_formatter)
2628 v.view()
2630 def print_errors(self) -> None:
2631 """Print tracebacks for every node with state "ERROR" in a Computation."""
2632 for n in self.nodes():
2633 if self.s[n] == States.ERROR:
2634 print(f"{n}")
2635 print("=" * len(str(n)))
2636 print()
2637 print(self.v[n].traceback)
2638 print()
2640 @classmethod
2641 def from_class(cls, definition_class: type, ignore_self: bool = True) -> "Computation":
2642 """Create a computation from a class with decorated methods."""
2643 comp = cls()
2644 obj = definition_class()
2645 populate_computation_from_class(comp, definition_class, obj, ignore_self=ignore_self)
2646 return comp
2648 @_notifies_subscribers()
2649 def inject_dependencies(self, dependencies: dict[Name, Any], *, force: bool = False) -> None:
2650 """Injects dependencies into the nodes of the current computation where nodes are in a placeholder state.
2652 (or all possible nodes when the 'force' parameter is set to True), using values
2653 provided in the 'dependencies' dictionary.
2655 Each key in the 'dependencies' dictionary corresponds to a node identifier, and the associated
2656 value is the dependency object to inject. If the value is a callable, it will be added as a calc node.
2658 :param dependencies: A dictionary where each key-value pair consists of a node identifier and
2659 its corresponding dependency object or a callable that returns the dependency object.
2660 :param force: A boolean flag that, when set to True, forces the replacement of existing node values
2661 with the ones provided in 'dependencies', regardless of their current state. Defaults to False.
2662 :return: None
2663 """
2664 for n in self.nodes():
2665 if force or self.s[n] == States.PLACEHOLDER:
2666 obj = dependencies.get(n)
2667 if obj is None:
2668 continue
2669 if callable(obj):
2670 self.add_node(n, obj)
2671 else:
2672 self.add_node(n, value=obj)