redroot 0.2.0__py3-none-any.whl

redroot/trace.py

@@ -0,0 +1,474 @@
+"""The :class:`Trace`: one recorded execution and its value graph."""
+
+from __future__ import annotations
+
+import dataclasses
+import decimal
+import sys
+import uuid
+from collections import deque
+from collections.abc import Callable, Mapping, Sequence
+from contextvars import Token
+from dataclasses import dataclass
+from types import MappingProxyType, TracebackType
+from typing import TYPE_CHECKING, Any, TypeVar
+
+from redroot._core import (
+    _WRAPPERS,
+    GUARD,
+    LEAF,
+    Node,
+    Traced,
+    _active,
+    active_trace,
+    add_node,
+    unwrap,
+)
+from redroot.ops import get_op
+from redroot.paths import child_key, format_key
+
+if TYPE_CHECKING:
+    from redroot.propagation import Propagation
+
+__all__ = ["Output", "Trace", "run", "track"]
+
+T = TypeVar("T")
+
+
+@dataclass(frozen=True, slots=True)
+class Output:
+    """A value the workflow produced, bound to the node that computed it.
+
+    Attributes:
+        key: Semantic id, e.g. ``"out:form1.line_12"``.
+        node: The node behind the value, or ``None`` if the value was not
+            traced ("not linked to the inputs").
+        value: The plain value at collection time.
+    """
+
+    key: str
+    node: Node | None
+    value: Any
+
+    @property
+    def linked(self) -> bool:
+        """Whether the value is linked to traced inputs."""
+        return self.node is not None
+
+
+class Trace:
+    """A recording of one execution, as a graph from inputs to outputs.
+
+    Use it as a context manager. Inside the block, operations on traced
+    values are recorded::
+
+        with Trace("run-42") as trace:
+            data = trace.track(extracted, root="ext")
+            result = workflow(data)
+            trace.collect(result, root="out")
+
+        trace.sources("out:form1.line_12")      # which inputs produced it
+        trace.propagate({"ext:a": 900})         # what changes if an input does
+
+    A trace is not thread-safe: record into it from one thread at a time.
+    """
+
+    def __init__(self, name: str | None = None, *, id: str | None = None) -> None:
+        self.id = id or uuid.uuid4().hex
+        self.name = name
+        self.decimal_context: decimal.Context = decimal.getcontext().copy()
+        """Decimal context used for re-evaluation; captured when recording ends."""
+        self._nodes: list[Node] = []
+        self._leaves: dict[str, Node] = {}
+        self._untracked: dict[str, Any] = {}
+        self._outputs: dict[str, Output] = {}
+        self._tokens: list[Token[Any]] = []
+        # decimal contexts nodes were computed under (see Node.ctx)
+        self._contexts: list[decimal.Context] = []
+        self._context_ids: dict[tuple[Any, ...], int] = {}
+
+    def __repr__(self) -> str:
+        name = f" {self.name!r}" if self.name else ""
+        return (
+            f"<Trace{name} nodes={len(self._nodes)} inputs={len(self._leaves)} "
+            f"outputs={len(self._outputs)} guards={len(self.guards)}>"
+        )
+
+    # ------------------------------------------------------------------
+    # Activation
+    # ------------------------------------------------------------------
+
+    def __enter__(self) -> Trace:
+        self._tokens.append(_active.set(self))
+        return self
+
+    def __exit__(
+        self,
+        exc_type: type[BaseException] | None,
+        exc: BaseException | None,
+        tb: TracebackType | None,
+    ) -> None:
+        _active.reset(self._tokens.pop())
+        self.decimal_context = decimal.getcontext().copy()
+
+    @property
+    def active(self) -> bool:
+        """Whether this trace is recording in the current context."""
+        return active_trace() is self
+
+    # ------------------------------------------------------------------
+    # Contents
+    # ------------------------------------------------------------------
+
+    @property
+    def nodes(self) -> Sequence[Node]:
+        """Every node, in recording (topological) order. Do not mutate."""
+        return self._nodes
+
+    @property
+    def decimal_contexts(self) -> Sequence[decimal.Context]:
+        """The ``decimal`` contexts that ``Decimal`` nodes were computed under."""
+        return tuple(self._contexts)
+
+    @property
+    def guards(self) -> list[Node]:
+        """Nodes whose values left tracking (comparisons, conversions...)."""
+        return [node for node in self._nodes if node.kind == GUARD]
+
+    @property
+    def inputs(self) -> Mapping[str, Node]:
+        """Keyed input (leaf) nodes."""
+        return MappingProxyType(self._leaves)
+
+    @property
+    def untracked_inputs(self) -> Mapping[str, Any]:
+        """Inputs passed to :meth:`track` that could not be traced (``None``, ``bool``...)."""
+        return MappingProxyType(self._untracked)
+
+    @property
+    def outputs(self) -> Mapping[str, Output]:
+        """Collected outputs by key."""
+        return MappingProxyType(self._outputs)
+
+    def output_values(self) -> dict[str, Any]:
+        """Plain value of every output, by key."""
+        return {key: out.value for key, out in self._outputs.items()}
+
+    # ------------------------------------------------------------------
+    # Inputs
+    # ------------------------------------------------------------------
+
+    def _add_leaf(self, value: Any, key: str | None, meta: dict[str, Any] | None) -> Node:
+        if key is not None:
+            self._check_new_key(key)
+        node = add_node(self, LEAF, (), None, value, LEAF, key=key, meta=meta)
+        if key is not None:
+            self._leaves[key] = node
+        return node
+
+    def _check_new_key(self, key: str) -> None:
+        if key in self._leaves or key in self._untracked or key in self._outputs:
+            raise ValueError(f"key {key!r} is already used in this trace")
+
+    def leaf(
+        self, value: Any, key: str | None = None, *, meta: dict[str, Any] | None = None
+    ) -> Any:
+        """Record ``value`` as an input and return its traced form.
+
+        Args:
+            value: An ``int``, ``float``, ``Decimal`` or ``str``.
+            key: Semantic id used to refer to the input, e.g. for edits.
+            meta: Free-form metadata stored on the node.
+        """
+        plain = unwrap(value)
+        factory = _WRAPPERS.get(type(plain))
+        if factory is None:
+            raise TypeError(f"cannot trace a value of type {type(plain).__name__}")
+        return factory(plain, self._add_leaf(plain, key, meta))
+
+    def track(self, data: T, root: str = "in") -> T:
+        """Return a copy of ``data`` whose scalar values are traced inputs.
+
+        Dicts, lists and tuples are copied; every ``int``, ``float``,
+        ``Decimal`` and ``str`` inside becomes an input keyed by its path,
+        e.g. ``"ext:bank.line_items[3].amount"``. Other values (``None``,
+        ``bool``, objects) are kept as-is and listed in
+        :attr:`untracked_inputs`: editing them requires re-execution.
+        """
+        return self._track(data, format_key(root))  # type: ignore[no-any-return]
+
+    def _track(self, value: Any, key: str) -> Any:
+        kind = type(value)
+        if kind is dict:
+            return {k: self._track(v, child_key(key, k)) for k, v in value.items()}
+        if kind is list:
+            return [self._track(v, child_key(key, i)) for i, v in enumerate(value)]
+        if kind is tuple:
+            return tuple(self._track(v, child_key(key, i)) for i, v in enumerate(value))
+        plain = unwrap(value)
+        factory = _WRAPPERS.get(type(plain))
+        if factory is not None:
+            return factory(plain, self._add_leaf(plain, key, None))
+        self._check_new_key(key)
+        self._untracked[key] = value
+        return value
+
+    # ------------------------------------------------------------------
+    # Outputs
+    # ------------------------------------------------------------------
+
+    def output(self, key: str, value: Any) -> Output:
+        """Record ``value`` as the output ``key``."""
+        self._check_new_key(key)
+        node = value._rt_node if isinstance(value, Traced) else None
+        if node is not None and node.trace is not self:
+            node = None
+        out = Output(key, node, unwrap(value))
+        self._outputs[key] = out
+        return out
+
+    def collect(self, data: Any, root: str = "out") -> None:
+        """Record every scalar value inside ``data`` as an output.
+
+        Walks mappings, lists, tuples, dataclasses and Pydantic models; each
+        leaf value becomes an output keyed by its path under ``root``.
+        """
+        self._collect(data, format_key(root))
+
+    def _collect(self, value: Any, key: str) -> None:
+        if isinstance(value, Traced):
+            self.output(key, value)
+        elif isinstance(value, Mapping):
+            for k, v in value.items():
+                self._collect(v, child_key(key, k))
+        elif isinstance(value, (list, tuple)):
+            for i, v in enumerate(value):
+                self._collect(v, child_key(key, i))
+        elif _is_pydantic_model(value):
+            for name in type(value).model_fields:
+                self._collect(getattr(value, name), child_key(key, name))
+        elif dataclasses.is_dataclass(value) and not isinstance(value, type):
+            for field in dataclasses.fields(value):
+                self._collect(getattr(value, field.name), child_key(key, field.name))
+        else:
+            self.output(key, value)
+
+    # ------------------------------------------------------------------
+    # Navigation
+    # ------------------------------------------------------------------
+
+    def node(self, ref: Any) -> Node:
+        """Resolve ``ref`` to one of this trace's nodes.
+
+        ``ref`` may be a :class:`Node`, a traced value, or the key of an
+        input or of a linked output.
+        """
+        if isinstance(ref, Node):
+            node = ref
+        elif isinstance(ref, Traced):
+            node = ref._rt_node
+        elif isinstance(ref, str):
+            if ref in self._leaves:
+                return self._leaves[ref]
+            out = self._outputs.get(ref)
+            if out is None:
+                raise KeyError(f"no input or output named {ref!r}")
+            if out.node is None:
+                raise LookupError(f"output {ref!r} is not linked to any traced input")
+            return out.node
+        else:
+            raise TypeError(f"cannot resolve a node from {type(ref).__name__}")
+        if node.trace is not self:
+            raise ValueError("node belongs to a different trace")
+        return node
+
+    def ancestors(self, ref: Any) -> list[Node]:
+        """``ref``'s node and every node it was computed from, in id order."""
+        start = self.node(ref)
+        seen: dict[int, Node] = {start.id: start}
+        queue = deque([start])
+        while queue:
+            for operand in queue.popleft().operands():
+                if operand.id not in seen:
+                    seen[operand.id] = operand
+                    queue.append(operand)
+        return [seen[i] for i in sorted(seen)]
+
+    def sources(self, ref: Any) -> list[Node]:
+        """The input nodes ``ref`` was computed from.
+
+        This answers "which extracted fields produced this value?".
+        """
+        return [node for node in self.ancestors(ref) if node.kind == LEAF]
+
+    def descendants(self, ref: Any) -> list[Node]:
+        """``ref``'s node and every node computed from it, in id order."""
+        start = self.node(ref)
+        reached = {start.id}
+        found = [start]
+        for node in self._nodes[start.id + 1 :]:
+            if any(operand.id in reached for operand in node.operands()):
+                reached.add(node.id)
+                found.append(node)
+        return found
+
+    def dependents(self, ref: Any) -> list[str]:
+        """Keys of the outputs computed from ``ref``.
+
+        This answers "which outputs use this extracted field?".
+        """
+        reached = {node.id for node in self.descendants(ref)}
+        return [
+            key
+            for key, out in self._outputs.items()
+            if out.node is not None and out.node.id in reached
+        ]
+
+    def explain(self, ref: Any, *, max_depth: int = 12) -> str:
+        """Render how ``ref`` was computed as an expression over the inputs.
+
+        >>> trace.explain("out:total")  # doctest: +SKIP
+        '(ext:a + ext:b) * 0.8'
+
+        Sub-expressions deeper than ``max_depth`` are shown as ``#<node id>``.
+        """
+        return _render(self.node(ref), max_depth, top=True)
+
+    # ------------------------------------------------------------------
+    # Propagation and serialization
+    # ------------------------------------------------------------------
+
+    def propagate(self, edits: Mapping[str, Any]) -> Propagation:
+        """Recompute every output after changing some inputs.
+
+        See :func:`redroot.propagation.propagate`.
+        """
+        from redroot.propagation import propagate
+
+        return propagate(self, edits)
+
+    def to_dict(self) -> dict[str, Any]:
+        """Serialize to JSON-compatible data. See :mod:`redroot.serialization`."""
+        from redroot.serialization import trace_to_dict
+
+        return trace_to_dict(self)
+
+    def to_json(self, **kwargs: Any) -> str:
+        """Serialize to a JSON string; ``kwargs`` go to :func:`json.dumps`."""
+        import json
+
+        return json.dumps(self.to_dict(), **kwargs)
+
+    @classmethod
+    def from_dict(cls, data: Mapping[str, Any]) -> Trace:
+        """Rebuild a trace from :meth:`to_dict` output."""
+        from redroot.serialization import trace_from_dict
+
+        return trace_from_dict(data)
+
+    @classmethod
+    def from_json(cls, text: str) -> Trace:
+        """Rebuild a trace from :meth:`to_json` output."""
+        import json
+
+        return cls.from_dict(json.loads(text))
+
+
+def _is_pydantic_model(value: Any) -> bool:
+    pydantic = sys.modules.get("pydantic")
+    return pydantic is not None and isinstance(value, pydantic.BaseModel)
+
+
+def _render(node: Node, depth: int, top: bool = False) -> str:
+    if node.kind == LEAF:
+        return node.key if node.key is not None else repr(node.value)
+    if depth <= 0:
+        return f"#{node.id}"
+    args = [_render_arg(arg, depth - 1) for arg in node.args]
+    args += [f"{k}={_render_arg(v, depth - 1)}" for k, v in (node.kwargs or {}).items()]
+    if node.op == "fstring":
+        return _render_fstring(node.args, depth - 1)
+    if node.op == "getslice" and len(args) == 4:
+        start, stop, step = (
+            "" if a is None else t for a, t in zip(node.args[1:], args[1:], strict=True)
+        )
+        return f"{args[0]}[{start}:{stop}{':' + step if step else ''}]"
+    spec = get_op(node.op)
+    style = spec.style if spec is not None else "call"
+    symbol = spec.symbol if spec is not None else None
+    text: str
+    if style == "infix" and len(args) == 2:
+        text = f"{args[0]} {symbol} {args[1]}"
+    elif style == "contains" and len(args) == 2:
+        text = f"{args[1]} in {args[0]}"
+    elif style == "prefix" and len(args) == 1:
+        return f"{symbol}{args[0]}"
+    elif style == "getitem" and len(args) == 2:
+        return f"{args[0]}[{args[1]}]"
+    elif style == "method" and args:
+        return f"{args[0]}.{symbol}({', '.join(args[1:])})"
+    else:
+        return f"{symbol or node.op}({', '.join(args)})"
+    return text if top else f"({text})"
+
+
+def _render_fstring(parts: tuple[Any, ...], depth: int) -> str:
+    out = []
+    for part in parts:
+        if isinstance(part, str):
+            out.append(part.replace("{", "{{").replace("}", "}}"))
+            continue
+        value, conversion, spec = part
+        conv = f"!{conversion}" if conversion else ""
+        spec_text = f":{_render_arg(spec, depth) if type(spec) is Node else spec}" if spec else ""
+        out.append(f"{{{_render_arg(value, depth)}{conv}{spec_text}}}")
+    return 'f"' + "".join(out).replace('"', '\\"') + '"'
+
+
+def _render_arg(arg: Any, depth: int) -> str:
+    kind = type(arg)
+    if kind is Node:
+        return _render(arg, depth)
+    if kind is list:
+        return "[" + ", ".join(_render_arg(a, depth) for a in arg) + "]"
+    if kind is tuple:
+        inner = ", ".join(_render_arg(a, depth) for a in arg)
+        return f"({inner},)" if len(arg) == 1 else f"({inner})"
+    if kind is dict:
+        return "{" + ", ".join(f"{k!r}: {_render_arg(v, depth)}" for k, v in arg.items()) + "}"
+    return repr(arg)
+
+
+# ----------------------------------------------------------------------
+# Module-level helpers
+# ----------------------------------------------------------------------
+
+
+def track(data: T, root: str = "in") -> T:
+    """Trace ``data`` as inputs of the active trace. See :meth:`Trace.track`."""
+    trace = active_trace()
+    if trace is None:
+        raise RuntimeError("redroot.track() needs an active Trace (use 'with Trace():')")
+    return trace.track(data, root)
+
+
+def run(
+    workflow: Callable[[Any], Any],
+    inputs: Any,
+    *,
+    input_root: str = "in",
+    output_root: str = "out",
+    name: str | None = None,
+) -> tuple[Trace, Any]:
+    """Run ``workflow(inputs)`` under a new trace.
+
+    Inputs are traced under ``input_root`` and every value in the returned
+    result is collected as an output under ``output_root``.
+
+    Returns:
+        The trace and the workflow's result.
+    """
+    with Trace(name) as trace:
+        result = workflow(trace.track(inputs, input_root))
+        trace.collect(result, output_root)
+    return trace, result