redroot 0.2.0__py3-none-any.whl

redroot/types.py

@@ -0,0 +1,385 @@
+"""Traced counterparts of ``int``, ``float``, ``Decimal`` and ``str``.
+
+A traced value *is* an instance of its builtin type (``isinstance(x, int)``
+holds, JSON encoders and C APIs accept it) and behaves exactly like it. While a
+:class:`~redroot.Trace` is active, every operation on it records a node;
+outside a trace, operations simply return plain builtin values.
+
+What each operation records:
+
+* arithmetic, ``round()``, ``math.floor()``, ``format()``, ``str()``, ``str``
+  and ``Decimal`` methods -> a re-evaluable node, returned as a traced value;
+* comparisons, truth tests, ``int()``/``float()``, ``hash()``, ``len()`` ->
+  a *guard*: the plain result is returned, and the trace remembers what it
+  was so that re-evaluation can tell whether it would change.
+
+Constructing ``TracedX(value, key=None, meta=None)`` converts ``value`` like
+``X(value)``. Inside an active trace, converting a value traced in that trace
+records a re-evaluable conversion; anything else becomes a new input, named
+``key`` if given. Outside a trace the value is detached: it behaves like a
+plain value and records nothing.
+"""
+
+from __future__ import annotations
+
+import math
+import operator
+from collections.abc import Callable
+from decimal import Decimal
+from typing import TYPE_CHECKING, Any
+
+from redroot._core import (
+    LEAF,
+    OP,
+    Node,
+    Observer,
+    Traced,
+    _active,
+    add_node,
+    binary,
+    call,
+    decimal_context_id,
+    observe_value,
+    record_raise,
+    register_wrapper,
+    unwrap,
+)
+from redroot.ops import (
+    BINARY_OPS,
+    COMPARISON_OPS,
+    DECIMAL_METHODS,
+    FLOAT_METHODS,
+    INT_METHODS,
+    STR_METHODS,
+    UNARY_OPS,
+    observe,
+    register_op,
+)
+
+if TYPE_CHECKING:
+    from pydantic import GetCoreSchemaHandler
+    from pydantic_core import CoreSchema
+
+__all__ = ["TracedDecimal", "TracedFloat", "TracedInt", "TracedStr"]
+
+
+def _construct(
+    cls: type[Any],
+    base: type[Any],
+    cast_op: str,
+    value: Any,
+    key: str | None,
+    meta: dict[str, Any] | None,
+) -> Any:
+    """Shared constructor logic.
+
+    ``TracedX(v)`` behaves like ``X(v)``. Inside an active trace, converting
+    a value traced in that trace records a re-evaluable cast; anything else
+    creates a new leaf (an input). Outside a trace the value is detached.
+    """
+    trace = _active.get()
+    source = value._rt_node if isinstance(value, Traced) else None
+    is_cast = trace is not None and key is None and source is not None and source.trace is trace
+    try:
+        plain = base(unwrap(value))
+    except Exception as exc:
+        if is_cast:
+            record_raise(trace, cast_op, (source,), None, exc)  # type: ignore[arg-type]  # is_cast implies a Trace
+        raise
+    if trace is None or isinstance(trace, Observer):
+        if trace is not None:
+            trace.touch_node(source)
+        node = Node(LEAF, (), None, plain, LEAF, key=key, meta=meta)
+    elif is_cast:
+        ctx = decimal_context_id(trace) if base is Decimal else None
+        node = add_node(trace, cast_op, (source,), None, plain, OP, meta=meta, ctx=ctx)
+    else:
+        node = trace._add_leaf(plain, key, meta)
+    obj: Any = base.__new__(cls, plain)
+    obj._rt_node = node
+    return obj
+
+
+def _factory(cls: type[Any], base: type[Any]) -> Callable[[Any, Node], Any]:
+    def wrap(value: Any, node: Node) -> Any:
+        obj: Any = base.__new__(cls, value)
+        obj._rt_node = node
+        return obj
+
+    return wrap
+
+
+class _TracedScalar(Traced):
+    """Behaviour shared by every traced type: repr, copying, pickling, Pydantic."""
+
+    __slots__ = ()
+    _base: type[Any]
+
+    def __repr__(self) -> str:
+        return repr(self._rt_node.value)
+
+    def __copy__(self) -> Any:
+        return self
+
+    def __deepcopy__(self, memo: dict[int, Any]) -> Any:
+        return self
+
+    def __reduce__(self) -> tuple[Any, ...]:
+        # Lineage does not survive serialization: unpickle as the plain value.
+        return (self._base, (self._rt_node.value,))
+
+    @classmethod
+    def __get_pydantic_core_schema__(
+        cls, source_type: Any, handler: GetCoreSchemaHandler
+    ) -> CoreSchema:
+        from pydantic_core import core_schema
+
+        schemas: dict[type, Callable[[], CoreSchema]] = {
+            int: core_schema.int_schema,
+            float: core_schema.float_schema,
+            Decimal: core_schema.decimal_schema,
+            str: core_schema.str_schema,
+        }
+        from_plain = core_schema.no_info_after_validator_function(cls, schemas[cls._base]())
+        return core_schema.json_or_python_schema(
+            json_schema=from_plain,
+            # Already-traced values pass through untouched, keeping their lineage.
+            python_schema=core_schema.union_schema(
+                [core_schema.is_instance_schema(cls), from_plain]
+            ),
+            serialization=core_schema.plain_serializer_function_ser_schema(unwrap),
+        )
+
+
+class TracedInt(_TracedScalar, int):
+    """An ``int`` that records how it was computed."""
+
+    # int/str are variable-size: instances keep the node in __dict__.
+    _base = int
+
+    def __new__(
+        cls, value: Any = 0, *, key: str | None = None, meta: dict[str, Any] | None = None
+    ) -> TracedInt:
+        """Convert ``value`` like ``int(value)`` (see the module docstring)."""
+        return _construct(cls, int, "int", value, key, meta)  # type: ignore[no-any-return]
+
+
+class TracedFloat(_TracedScalar, float):
+    """A ``float`` that records how it was computed."""
+
+    __slots__ = ("_rt_node",)
+    _base = float
+
+    def __new__(
+        cls, value: Any = 0.0, *, key: str | None = None, meta: dict[str, Any] | None = None
+    ) -> TracedFloat:
+        """Convert ``value`` like ``float(value)`` (see the module docstring)."""
+        return _construct(cls, float, "float", value, key, meta)  # type: ignore[no-any-return]
+
+
+class TracedDecimal(_TracedScalar, Decimal):
+    """A ``decimal.Decimal`` that records how it was computed."""
+
+    __slots__ = ("_rt_node",)
+    _base = Decimal
+
+    def __new__(
+        cls, value: Any = "0", *, key: str | None = None, meta: dict[str, Any] | None = None
+    ) -> TracedDecimal:
+        """Convert ``value`` like ``Decimal(value)`` (see the module docstring)."""
+        return _construct(cls, Decimal, "decimal", value, key, meta)  # type: ignore[no-any-return]
+
+
+class TracedStr(_TracedScalar, str):
+    """A ``str`` that records how it was computed."""
+
+    # int/str are variable-size: instances keep the node in __dict__.
+    _base = str
+
+    def __new__(
+        cls, value: Any = "", *, key: str | None = None, meta: dict[str, Any] | None = None
+    ) -> TracedStr:
+        """Convert ``value`` like ``str(value)`` (see the module docstring)."""
+        return _construct(cls, str, "str", value, key, meta)  # type: ignore[no-any-return]
+
+
+# --------------------------------------------------------------------------
+# Method installation. Methods are attached after class creation so that
+# static type checkers keep seeing the builtin signatures: a TracedInt is
+# typed as an int, which is exactly what it behaves like.
+# --------------------------------------------------------------------------
+
+
+def _binary_pair(op: str) -> tuple[Callable[..., Any], Callable[..., Any]]:
+    fn = BINARY_OPS[op][0]
+
+    def forward(self: Any, other: Any) -> Any:
+        return binary(op, fn, self, other)
+
+    def reflected(self: Any, other: Any) -> Any:
+        return binary(op, fn, other, self)
+
+    return forward, reflected
+
+
+def _pow(self: Any, other: Any, modulo: Any = None) -> Any:
+    if modulo is None:
+        return binary("pow", pow, self, other)
+    return call("pow", pow, (self, other, modulo))
+
+
+def _rpow(self: Any, other: Any, modulo: Any = None) -> Any:
+    if modulo is None:
+        return binary("pow", pow, other, self)
+    return call("pow", pow, (other, self, modulo))
+
+
+def _comparison(op: str) -> Callable[..., Any]:
+    fn = COMPARISON_OPS[op][0]
+
+    def compare(self: Any, other: Any) -> Any:
+        return observe_value(op, fn, self, other)
+
+    return compare
+
+
+def _guard(op: str, fn: Callable[[Any], Any]) -> Callable[..., Any]:
+    def method(self: Any) -> Any:
+        return observe_value(op, fn, self)
+
+    return method
+
+
+def _hash(self: Any) -> int:
+    # Guard the whole value, not just its hash: -1 and -2 hash alike, yet a
+    # dict keyed by them differs.
+    return hash(observe_value("observe", observe, self))
+
+
+def _function(op: str, fn: Callable[..., Any]) -> Callable[..., Any]:
+    def method(self: Any, *args: Any) -> Any:
+        return call(op, fn, (self, *args))
+
+    return method
+
+
+def _round(self: Any, ndigits: Any = None) -> Any:
+    if ndigits is None:
+        return call("round", round, (self,))
+    return call("round", round, (self, ndigits))
+
+
+def _method(prefix: str, name: str, owner: type[Any]) -> Callable[..., Any]:
+    op = f"{prefix}.{name}"
+    fn = getattr(owner, name)
+
+    def method(self: Any, *args: Any, **kwargs: Any) -> Any:
+        return call(op, fn, (self, *args), kwargs or None)
+
+    method.__name__ = name
+    method.__doc__ = fn.__doc__
+    return method
+
+
+def _install(cls: type[Any], name: str, func: Callable[..., Any]) -> None:
+    if not name.startswith("__"):
+        func.__qualname__ = f"{cls.__name__}.{name}"
+    setattr(cls, name, func)
+
+
+def _install_number(
+    cls: type[Any], binary_ops: tuple[str, ...], unary_ops: tuple[str, ...]
+) -> None:
+    for op in binary_ops:
+        forward, reflected = _binary_pair(op)
+        _install(cls, f"__{op}__", forward)
+        _install(cls, f"__r{op}__", reflected)
+    _install(cls, "__pow__", _pow)
+    _install(cls, "__rpow__", _rpow)
+    for op in unary_ops:
+        _install(cls, f"__{op}__", _function(op, UNARY_OPS[op][0]))
+    for op in COMPARISON_OPS:
+        _install(cls, f"__{op}__", _comparison(op))
+    _install(cls, "__abs__", _function("abs", abs))
+    _install(cls, "__round__", _round)
+    _install(cls, "__floor__", _function("floor", math.floor))
+    _install(cls, "__ceil__", _function("ceil", math.ceil))
+    _install(cls, "__trunc__", _function("trunc", math.trunc))
+    _install(cls, "__divmod__", _function("divmod", divmod))
+    _install(cls, "__rdivmod__", lambda self, other: call("divmod", divmod, (other, self)))
+    _install(cls, "__format__", _function("format", format))
+    _install(cls, "__str__", _function("str", str))
+    _install(cls, "__bool__", _guard("bool", bool))
+    _install(cls, "__int__", _guard("int", int))
+    _install(cls, "__float__", _guard("float", float))
+    _install(cls, "__complex__", _guard("complex", complex))
+    _install(cls, "__hash__", _hash)
+
+
+_ARITHMETIC = ("add", "sub", "mul", "truediv", "floordiv", "mod")
+
+_install_number(
+    TracedInt, (*_ARITHMETIC, "lshift", "rshift", "and", "or", "xor"), ("neg", "pos", "invert")
+)
+_install(TracedInt, "__index__", _guard("index", operator.index))
+_install_number(TracedFloat, _ARITHMETIC, ("neg", "pos"))
+_install_number(TracedDecimal, _ARITHMETIC, ("neg", "pos"))
+
+for _prefix, _cls, _owner, _names in (
+    ("int", TracedInt, int, INT_METHODS),
+    ("float", TracedFloat, float, FLOAT_METHODS),
+    ("decimal", TracedDecimal, Decimal, DECIMAL_METHODS),
+    ("str", TracedStr, str, STR_METHODS),
+):
+    for _name in _names:
+        _install(_cls, _name, _method(_prefix, _name, _owner))
+
+
+def _getslice(value: Any, start: Any, stop: Any, step: Any) -> Any:
+    return value[start:stop:step]
+
+
+register_op("getslice", _getslice)
+
+
+def _str_getitem(self: Any, key: Any) -> Any:
+    if isinstance(key, slice):
+        return call("getslice", _getslice, (self, key.start, key.stop, key.step))
+    return call("getitem", operator.getitem, (self, key))
+
+
+def _str_iter(self: Any) -> Any:
+    # Iterating exposes every character: treat the whole string as observed.
+    return iter(observe_value("observe", lambda v: v, self))
+
+
+def _str_join(self: Any, iterable: Any) -> Any:
+    return call("str.join", str.join, (self, list(iterable)))
+
+
+for _op in ("add", "mul"):
+    _forward, _reflected = _binary_pair(_op)
+    _install(TracedStr, f"__{_op}__", _forward)
+    _install(TracedStr, f"__r{_op}__", _reflected)
+_install(TracedStr, "__mod__", lambda self, other: call("mod", operator.mod, (self, other)))
+_install(TracedStr, "__rmod__", lambda self, other: call("mod", operator.mod, (other, self)))
+for _op in COMPARISON_OPS:
+    _install(TracedStr, f"__{_op}__", _comparison(_op))
+_install(TracedStr, "__getitem__", _str_getitem)
+_install(
+    TracedStr,
+    "__contains__",
+    lambda self, item: observe_value("contains", operator.contains, self, item),
+)
+_install(TracedStr, "__len__", _guard("len", len))
+_install(TracedStr, "__bool__", _guard("bool", bool))
+_install(TracedStr, "__hash__", _hash)
+_install(TracedStr, "__iter__", _str_iter)
+_install(TracedStr, "__format__", _function("format", format))
+_install(TracedStr, "__str__", lambda self: self)
+_install(TracedStr, "join", _str_join)
+
+register_wrapper(int, _factory(TracedInt, int))
+register_wrapper(float, _factory(TracedFloat, float))
+register_wrapper(Decimal, _factory(TracedDecimal, Decimal))
+register_wrapper(str, _factory(TracedStr, str))

redroot/validation.py

@@ -0,0 +1,221 @@
+"""Checks that tell you whether tracing is trustworthy for a workflow.
+
+Run these on real recorded workflows before relying on propagation:
+
+* :func:`coverage` - how many outputs are linked to inputs, how many can be
+  recomputed, how many inputs feed guards.
+* :func:`check_identity` - re-evaluating every node from its recorded
+  operands reproduces the recorded value exactly.
+* :func:`check_perturbation` - after an edit, propagated outputs agree with a
+  real re-execution wherever propagation claimed to know the answer.
+"""
+
+from __future__ import annotations
+
+import decimal
+from collections import deque
+from collections.abc import Callable, Mapping
+from dataclasses import dataclass
+from typing import Any
+
+from redroot._core import GUARD, LEAF, Node, same_value
+from redroot.paths import apply_edits
+from redroot.propagation import Propagation, Verification, evaluate_outcome, node_function
+from redroot.trace import Trace, run
+
+__all__ = [
+    "Coverage",
+    "IdentityMismatch",
+    "IdentityReport",
+    "PerturbationReport",
+    "check_identity",
+    "check_perturbation",
+    "coverage",
+]
+
+
+@dataclass(frozen=True, slots=True)
+class Coverage:
+    """How much of a trace's output the graph can explain and recompute."""
+
+    outputs: int
+    linked: int
+    """Outputs computed from at least one keyed input."""
+    replayable: int
+    """Linked outputs whose every ancestor operation can be re-evaluated."""
+    inputs: int
+    untracked_inputs: int
+    guards: int
+    guarded_inputs: tuple[str, ...]
+    """Input keys that feed at least one guard: edits to them may force re-execution."""
+
+    @property
+    def unlinked(self) -> int:
+        """Outputs not backed by a traced value."""
+        return self.outputs - self.linked
+
+    @property
+    def linked_ratio(self) -> float:
+        """Share of outputs that are linked (1.0 if there are none)."""
+        return self.linked / self.outputs if self.outputs else 1.0
+
+    @property
+    def replayable_ratio(self) -> float:
+        """Share of outputs that can be recomputed (1.0 if there are none)."""
+        return self.replayable / self.outputs if self.outputs else 1.0
+
+    def __str__(self) -> str:
+        return "\n".join(
+            [
+                f"outputs ........... {self.outputs}",
+                f"  linked .......... {self.linked} ({self.linked_ratio:.0%})",
+                f"  replayable ...... {self.replayable} ({self.replayable_ratio:.0%})",
+                f"inputs ............ {self.inputs} (+{self.untracked_inputs} untracked)",
+                f"  feeding guards .. {len(self.guarded_inputs)}",
+                f"guards ............ {self.guards}",
+            ]
+        )
+
+
+def coverage(trace: Trace) -> Coverage:
+    """Measure how well ``trace`` covers its outputs."""
+    replayable: dict[int, bool] = {}
+    keyed: dict[int, bool] = {}  # computed from at least one keyed input
+    for node in trace.nodes:
+        if node.kind == LEAF:
+            replayable[node.id] = True
+            keyed[node.id] = node.key is not None
+            continue
+        operands = list(node.operands())
+        replayable[node.id] = node_function(node) is not None and all(
+            replayable[operand.id] for operand in operands
+        )
+        keyed[node.id] = any(keyed[operand.id] for operand in operands)
+
+    # Values built only from anonymous leaves (e.g. rebuilt inside a library
+    # with type(x)(...)) are not linked to the inputs, whatever their type.
+    linked = [
+        out.node for out in trace.outputs.values() if out.node is not None and keyed[out.node.id]
+    ]
+    return Coverage(
+        outputs=len(trace.outputs),
+        linked=len(linked),
+        replayable=sum(1 for node in linked if replayable[node.id]),
+        inputs=len(trace.inputs),
+        untracked_inputs=len(trace.untracked_inputs),
+        guards=sum(1 for node in trace.nodes if node.kind == GUARD),
+        guarded_inputs=_guarded_inputs(trace),
+    )
+
+
+def _guarded_inputs(trace: Trace) -> tuple[str, ...]:
+    seen: set[int] = set()
+    queue = deque(node for node in trace.nodes if node.kind == GUARD)
+    keys: list[tuple[int, str]] = []
+    while queue:
+        node = queue.popleft()
+        for operand in node.operands():
+            if operand.id in seen:
+                continue
+            seen.add(operand.id)
+            if operand.kind == LEAF:
+                if operand.key is not None:
+                    keys.append((operand.id, operand.key))
+            else:
+                queue.append(operand)
+    return tuple(key for _, key in sorted(keys))
+
+
+@dataclass(frozen=True, slots=True)
+class IdentityMismatch:
+    """A node whose re-evaluation disagrees with its recorded value."""
+
+    node: Node
+    recorded: Any
+    recomputed: Any
+    error: str | None = None
+
+
+@dataclass(frozen=True, slots=True)
+class IdentityReport:
+    """Result of :func:`check_identity`."""
+
+    checked: int
+    skipped: int
+    """Nodes that are not replayable (e.g. LLM calls)."""
+    mismatches: tuple[IdentityMismatch, ...]
+
+    @property
+    def ok(self) -> bool:
+        """Whether every node re-evaluated to its recorded value."""
+        return not self.mismatches
+
+
+def check_identity(trace: Trace) -> IdentityReport:
+    """Re-evaluate every node from its recorded operands and compare.
+
+    Any mismatch means an operation did not behave like its registered
+    implementation (or is not deterministic), so propagation through it
+    cannot be trusted.
+    """
+    checked = skipped = 0
+    mismatches: list[IdentityMismatch] = []
+    with decimal.localcontext(trace.decimal_context):
+        for node in trace.nodes:
+            if node.kind == LEAF:
+                continue
+            if node_function(node) is None:
+                skipped += 1
+                continue
+            checked += 1
+            try:
+                recomputed = evaluate_outcome(node)
+            except Exception as exc:
+                mismatches.append(
+                    IdentityMismatch(node, node.value, None, f"{type(exc).__name__}: {exc}")
+                )
+                continue
+            if not same_value(recomputed, node.value):
+                mismatches.append(IdentityMismatch(node, node.value, recomputed))
+    return IdentityReport(checked, skipped, tuple(mismatches))
+
+
+@dataclass(frozen=True, slots=True)
+class PerturbationReport:
+    """Result of :func:`check_perturbation`."""
+
+    trace: Trace
+    """The trace of the original run."""
+    propagation: Propagation
+    """What the graph predicted for the edit."""
+    retrace: Trace
+    """The trace of the re-execution on edited inputs."""
+    verification: Verification
+
+    @property
+    def consistent(self) -> bool:
+        """Whether every exact prediction matched the re-execution."""
+        return self.verification.ok
+
+
+def check_perturbation(
+    workflow: Callable[[Any], Any],
+    inputs: Any,
+    edits: Mapping[str, Any],
+    *,
+    input_root: str = "in",
+    output_root: str = "out",
+) -> PerturbationReport:
+    """Edit inputs, then compare graph propagation with a real re-execution.
+
+    The workflow runs twice under tracing: on ``inputs``, then on ``inputs``
+    with ``edits`` applied. Outputs propagation reports as ``UPDATED``,
+    ``UNCHANGED`` or ``UNLINKED`` must match the re-execution; ``STALE``
+    outputs are expected to need it.
+    """
+    trace, _ = run(workflow, inputs, input_root=input_root, output_root=output_root)
+    propagation = trace.propagate(edits)
+    edited = apply_edits(inputs, edits, root=input_root)
+    retrace, _ = run(workflow, edited, input_root=input_root, output_root=output_root)
+    verification = propagation.verify(retrace.output_values())
+    return PerturbationReport(trace, propagation, retrace, verification)

redroot/visualizer/__init__.py

@@ -0,0 +1,11 @@
+"""Exports and viewers for traces.
+
+``to_networkx`` and ``to_graphviz`` need ``redroot[viz]``; the web viewer
+(``redroot visualize trace.json``) needs ``redroot[web]``.
+"""
+
+from redroot.visualizer.data import graph_data
+from redroot.visualizer.graphviz import to_graphviz
+from redroot.visualizer.networkx import to_networkx
+
+__all__ = ["graph_data", "to_graphviz", "to_networkx"]

redroot/visualizer/data.py

@@ -0,0 +1,63 @@
+"""A display-oriented view of a trace, shared by the exporters and the web viewer."""
+
+from __future__ import annotations
+
+from typing import Any
+
+from redroot._core import Node
+from redroot.trace import Trace
+
+
+def _short(value: Any, limit: int) -> str:
+    text = repr(value)
+    return text if len(text) <= limit else text[: limit - 1] + "…"
+
+
+def node_label(node: Node, limit: int = 24) -> str:
+    """A one-line label: the key for inputs, ``op = value`` otherwise."""
+    if node.key is not None:
+        return node.key
+    return f"{node.op} = {_short(node.value, limit)}"
+
+
+def graph_data(trace: Trace, *, label_limit: int = 24) -> dict[str, Any]:
+    """Return nodes, edges and outputs of ``trace`` as JSON-compatible data.
+
+    Edges point from an operand to the node computed from it.
+    """
+    nodes = [
+        {
+            "id": node.id,
+            "op": node.op,
+            "kind": node.kind,
+            "key": node.key,
+            "label": node_label(node, label_limit),
+            "value": _short(node.value, 500),
+            "meta": {k: _short(v, 200) for k, v in (node.meta or {}).items()},
+        }
+        for node in trace.nodes
+    ]
+    edges = [
+        {"source": operand_id, "target": node.id}
+        for node in trace.nodes
+        for operand_id in dict.fromkeys(operand.id for operand in node.operands())
+    ]
+    outputs = []
+    for key, out in trace.outputs.items():
+        linked = out.node is not None
+        outputs.append(
+            {
+                "key": key,
+                "node": out.node.id if out.node is not None else None,
+                "value": _short(out.value, 500),
+                "expression": trace.explain(key) if linked else None,
+                "sources": [s.key or f"#{s.id}" for s in trace.sources(key)] if linked else [],
+            }
+        )
+    return {
+        "id": trace.id,
+        "name": trace.name,
+        "nodes": nodes,
+        "edges": edges,
+        "outputs": outputs,
+    }