redroot 0.2.0__py3-none-any.whl

redroot/__init__.py

@@ -0,0 +1,71 @@
+"""RedRoot: value-level lineage for Python, in both directions.
+
+Record a computation once, then ask where any output came from, and what
+every output becomes when an input changes, without running it again::
+
+    import redroot as rr
+
+    with rr.Trace() as trace:
+        data = trace.track({"a": 1200, "b": 300}, root="ext")
+        trace.collect({"total": data["a"] + data["b"]}, root="out")
+
+    trace.sources("out:total")                   # [ext:a, ext:b]
+    trace.propagate({"ext:a": 900}).outputs      # out:total -> 1200, exactly
+
+See the README for the full guide.
+"""
+
+from redroot._core import Node, Traced, active_trace, node_of, unwrap
+from redroot.functions import annotate, derive, traced, traced_llm
+from redroot.paths import apply_edits, format_key, parse_key
+from redroot.propagation import (
+    Mismatch,
+    OutputChange,
+    Propagation,
+    Reason,
+    ReasonKind,
+    Status,
+    Verification,
+)
+from redroot.trace import Output, Trace, run, track
+from redroot.types import TracedDecimal, TracedFloat, TracedInt, TracedStr
+
+__version__ = "0.2.0"
+
+
+def is_traced(value: object) -> bool:
+    """Whether ``value`` is a traced value."""
+    return isinstance(value, Traced)
+
+
+__all__ = [
+    "Mismatch",
+    "Node",
+    "Output",
+    "OutputChange",
+    "Propagation",
+    "Reason",
+    "ReasonKind",
+    "Status",
+    "Trace",
+    "Traced",
+    "TracedDecimal",
+    "TracedFloat",
+    "TracedInt",
+    "TracedStr",
+    "Verification",
+    "__version__",
+    "active_trace",
+    "annotate",
+    "apply_edits",
+    "derive",
+    "format_key",
+    "is_traced",
+    "node_of",
+    "parse_key",
+    "run",
+    "traced",
+    "traced_llm",
+    "track",
+    "unwrap",
+]

redroot/_core.py

@@ -0,0 +1,556 @@
+"""Recording machinery shared by the traced types and :class:`~redroot.Trace`.
+
+This module is internal. It holds the hot paths that run on every traced
+operation, so it favours plain functions and ``__slots__`` over abstraction.
+
+Invariants (see ``AGENTS.md``):
+
+* ``Node.value`` is always a *plain* value, never a traced one.
+* ``Node.args`` hold, in call order, either a :class:`Node` (a traced operand
+  recorded in the same trace) or a plain constant. Nothing is dropped.
+* Library code never calls ``int()``/``float()``/``str()``/``hash()``/``==``
+  on a traced value: those are recorded as guards. Use :func:`unwrap`.
+"""
+
+from __future__ import annotations
+
+import decimal
+import math
+from collections.abc import Callable, Iterator
+from contextvars import ContextVar
+from decimal import Decimal
+from typing import TYPE_CHECKING, Any, Final
+
+from redroot.ops import shape
+
+if TYPE_CHECKING:
+    from redroot.trace import Trace
+
+LEAF: Final = "leaf"
+"""Kind of an input value."""
+OP: Final = "op"
+"""Kind of a value computed from other values."""
+GUARD: Final = "guard"
+"""Kind of a value that left tracking (a comparison outcome, an ``int()``...).
+
+Execution may have depended on it in ways the graph cannot see, so if it
+changes on re-evaluation the graph can no longer be trusted.
+"""
+
+
+class Node:
+    """One recorded value in a :class:`~redroot.Trace`.
+
+    Attributes:
+        id: Position in the trace. Operands always have smaller ids, so id
+            order is a topological order.
+        op: Name of the operation (see :mod:`redroot.ops`); ``"leaf"`` for
+            inputs.
+        args: Operands in call order: :class:`Node` references or constants.
+            May nest inside lists, tuples and dicts.
+        kwargs: Keyword operands, or ``None``.
+        value: The plain value computed when the trace was recorded.
+        kind: ``"leaf"``, ``"op"`` or ``"guard"``.
+        key: Semantic id of a leaf, e.g. ``"ext:bank.line_items[3].amount"``.
+        meta: Free-form metadata (model name, source document...), or ``None``.
+        trace: The trace that owns the node; ``None`` if detached.
+        ctx: Index of the ``decimal`` context the node was computed under, in
+            the trace's contexts, or ``None`` if no ``Decimal`` was involved.
+        fn: The function that computed the node, when it is not the
+            registered one (``@traced`` functions), or :data:`NOT_REPLAYABLE`.
+    """
+
+    __slots__ = ("args", "ctx", "fn", "id", "key", "kind", "kwargs", "meta", "op", "trace", "value")
+
+    def __init__(
+        self,
+        op: str,
+        args: tuple[Any, ...],
+        kwargs: dict[str, Any] | None,
+        value: Any,
+        kind: str,
+        *,
+        key: str | None = None,
+        meta: dict[str, Any] | None = None,
+    ) -> None:
+        self.id = -1
+        self.op = op
+        self.args = args
+        self.kwargs = kwargs
+        self.value = value
+        self.kind = kind
+        self.key = key
+        self.meta = meta
+        self.trace: Trace | None = None
+        self.ctx: int | None = None
+        self.fn: Callable[..., Any] | None = None
+
+    def __repr__(self) -> str:
+        label = self.key if self.key is not None else self.op
+        return f"<Node #{self.id} {self.kind} {label}={self.value!r}>"
+
+    @property
+    def is_leaf(self) -> bool:
+        return self.kind == LEAF
+
+    @property
+    def is_guard(self) -> bool:
+        return self.kind == GUARD
+
+    def operands(self) -> Iterator[Node]:
+        """Yield the nodes this node was computed from, in argument order."""
+        yield from iter_nodes(self.args)
+        if self.kwargs:
+            yield from iter_nodes(self.kwargs)
+
+
+def iter_nodes(obj: Any) -> Iterator[Node]:
+    """Yield every :class:`Node` referenced in a (possibly nested) argument."""
+    kind = type(obj)
+    if kind is Node:
+        yield obj
+    elif kind is list or isinstance(obj, tuple):
+        for item in obj:
+            yield from iter_nodes(item)
+    elif kind is dict:
+        for item in obj.values():
+            yield from iter_nodes(item)
+
+
+def rebuild_tuple(template: tuple[Any, ...], items: Iterator[Any] | list[Any]) -> tuple[Any, ...]:
+    """Build a tuple of ``template``'s type (plain or named) from ``items``."""
+    if type(template) is tuple:
+        return tuple(items)
+    return type(template)._make(items)  # type: ignore[attr-defined, no-any-return]
+
+
+class Raised:
+    """The recorded outcome of an operation that raised an exception.
+
+    Operations on traced values that raise are recorded as guards with this
+    value, so code that catches the exception (``try``/``except``) is guarded:
+    if an edit makes the operation stop raising, or raise something else,
+    the guard flips.
+    """
+
+    __slots__ = ("type_name",)
+
+    def __init__(self, type_name: str) -> None:
+        self.type_name = type_name
+
+    @classmethod
+    def of(cls, exc: BaseException) -> Raised:
+        kind = type(exc)
+        module = "" if kind.__module__ == "builtins" else f"{kind.__module__}."
+        return cls(module + kind.__qualname__)
+
+    def __eq__(self, other: object) -> bool:
+        return isinstance(other, Raised) and other.type_name == self.type_name
+
+    def __hash__(self) -> int:
+        return hash(self.type_name)
+
+    def __repr__(self) -> str:
+        return f"<raised {self.type_name}>"
+
+
+def _not_replayable(*_args: Any, **_kwargs: Any) -> Any:
+    raise LookupError("this operation cannot be re-evaluated")
+
+
+NOT_REPLAYABLE: Final = _not_replayable
+"""Marks a node (``Node.fn``) whose operation must never be re-run."""
+
+
+# --------------------------------------------------------------------------
+# Active trace
+# --------------------------------------------------------------------------
+
+
+class Observer:
+    """Stands in for the active trace while a ``@traced`` function runs.
+
+    Nothing is recorded, but every node of the trace whose value the function
+    touches is collected, so the call's true inputs are known even when they
+    reach it through closures, globals, sets or dict views.
+    """
+
+    __slots__ = ("touched", "trace")
+
+    def __init__(self, trace: Trace) -> None:
+        self.trace = trace
+        self.touched: dict[int, Node] = {}
+
+    def touch_node(self, node: Node | None) -> None:
+        if node is not None and node.trace is self.trace:
+            self.touched[node.id] = node
+
+    def touch(self, *values: Any) -> None:
+        for value in values:
+            if isinstance(value, Traced):
+                node = value._rt_node
+                if node.trace is self.trace:
+                    self.touched[node.id] = node
+            elif type(value) in (list, tuple, dict):
+                self.touch(*(value.values() if type(value) is dict else value))
+
+
+_active: ContextVar[Trace | Observer | None] = ContextVar("redroot_active_trace", default=None)
+
+
+def active_trace() -> Trace | None:
+    """Return the trace recording in the current context, if any."""
+    trace = _active.get()
+    return None if type(trace) is Observer else trace  # type: ignore[return-value]
+
+
+# --------------------------------------------------------------------------
+# Traced values
+# --------------------------------------------------------------------------
+
+
+class Traced:
+    """Base class of every traced value type.
+
+    Each instance carries the :class:`Node` describing how it was produced.
+    """
+
+    __slots__ = ()
+    _rt_node: Node
+    _base: type[Any]
+    """The builtin type this traced type stands for."""
+
+
+def unwrap(value: Any) -> Any:
+    """Return the plain value behind ``value`` (``value`` itself if untraced).
+
+    Never triggers guards, so it is safe to use anywhere.
+    """
+    return value._rt_node.value if isinstance(value, Traced) else value
+
+
+def node_of(value: Any) -> Node | None:
+    """Return the :class:`Node` behind a traced value, or ``None``."""
+    return value._rt_node if isinstance(value, Traced) else None
+
+
+def same_value(a: Any, b: Any) -> bool:
+    """Whether two plain values are interchangeable as results.
+
+    Stricter than ``==``: types must match (``1 != 1.0 != True``), ``-0.0``
+    differs from ``0.0``, ``Decimal("1.0")`` differs from ``Decimal("1.00")``
+    (they print differently), and NaN equals NaN.
+    """
+    kind = type(a)
+    if kind is not type(b):
+        return False
+    if kind is float:
+        if a != a:  # NaN
+            return bool(b != b)
+        return bool(a == b and math.copysign(1.0, a) == math.copysign(1.0, b))
+    if kind is Decimal:
+        return bool(a.as_tuple() == b.as_tuple())
+    if kind is list or kind is tuple or isinstance(a, tuple):
+        return len(a) == len(b) and all(same_value(x, y) for x, y in zip(a, b, strict=True))
+    if kind is dict:
+        return list(a) == list(b) and all(same_value(a[k], b[k]) for k in a)
+    try:
+        result = a == b
+    except Exception:
+        return False
+    return result if isinstance(result, bool) else False
+
+
+# Exact plain type -> factory building the traced counterpart around a node.
+_WRAPPERS: dict[type, Callable[[Any, Node], Any]] = {}
+
+
+def register_wrapper(plain_type: type, factory: Callable[[Any, Node], Any]) -> None:
+    """Declare which traced type wraps results of ``plain_type``."""
+    _WRAPPERS[plain_type] = factory
+
+
+# --------------------------------------------------------------------------
+# Recording
+# --------------------------------------------------------------------------
+
+
+def add_node(
+    trace: Trace,
+    op: str,
+    args: tuple[Any, ...],
+    kwargs: dict[str, Any] | None,
+    value: Any,
+    kind: str,
+    *,
+    key: str | None = None,
+    meta: dict[str, Any] | None = None,
+    ctx: int | None = None,
+    fn: Callable[..., Any] | None = None,
+) -> Node:
+    """Append a node to ``trace`` and return it."""
+    node = Node(op, args, kwargs, value, kind, key=key, meta=meta)
+    nodes = trace._nodes
+    node.id = len(nodes)
+    node.trace = trace
+    node.ctx = ctx
+    node.fn = fn
+    nodes.append(node)
+    return node
+
+
+def decimal_context_id(trace: Trace) -> int:
+    """Intern the current ``decimal`` context in ``trace`` and return its index."""
+    c = decimal.getcontext()
+    fingerprint = (id(c), c.prec, c.rounding, c.Emin, c.Emax, c.capitals, c.clamp)
+    index = trace._context_ids.get(fingerprint)
+    if index is None:
+        index = trace._context_ids[fingerprint] = len(trace._contexts)
+        trace._contexts.append(c.copy())
+    return index
+
+
+def _uses_decimal(values: Any) -> bool:
+    return any(type(v) is Decimal for v in values)
+
+
+def _is_namedtuple(value: Any) -> bool:
+    return isinstance(value, tuple) and hasattr(type(value), "_fields")
+
+
+def is_projectable(value: Any) -> bool:
+    """Whether ``value`` is a container whose items are traced individually."""
+    kind = type(value)
+    return kind is list or kind is tuple or kind is dict or _is_namedtuple(value)
+
+
+def record_result(
+    trace: Trace,
+    op: str,
+    args: tuple[Any, ...],
+    kwargs: dict[str, Any] | None,
+    value: Any,
+    meta: dict[str, Any] | None = None,
+    ctx: int | None = None,
+    fn: Callable[..., Any] | None = None,
+) -> Any:
+    """Record ``value`` as the result of ``op(*args, **kwargs)``.
+
+    Returns the traced form of ``value``:
+
+    * a scalar with a traced counterpart (``int``, ``float``, ``Decimal``,
+      ``str``) is wrapped;
+    * a ``list``, ``tuple`` or ``dict`` is rebuilt from per-item projections,
+      plus a guard on its shape (length or keys);
+    * anything else (``bool``, ``None``, objects) is returned unchanged and
+      recorded as a guard, since it leaves tracking.
+    """
+    kind = type(value)
+    if isinstance(value, Traced) or kind is list or kind is dict or isinstance(value, tuple):
+        value = unwrap_deep(value)  # e.g. a function returned one of its traced arguments
+    factory = _WRAPPERS.get(type(value))
+    if factory is not None:
+        return factory(
+            value, add_node(trace, op, args, kwargs, value, OP, meta=meta, ctx=ctx, fn=fn)
+        )
+    if is_projectable(value):
+        node = add_node(trace, op, args, kwargs, value, OP, meta=meta, ctx=ctx, fn=fn)
+        return project(trace, node)
+    add_node(trace, op, args, kwargs, value, GUARD, meta=meta, ctx=ctx, fn=fn)
+    return value
+
+
+def record_raise(
+    trace: Trace,
+    op: str,
+    args: tuple[Any, ...],
+    kwargs: dict[str, Any] | None,
+    exc: BaseException,
+    fn: Callable[..., Any] | None = None,
+) -> None:
+    """Record that ``op(*args, **kwargs)`` raised ``exc``, as a guard."""
+    ctx = decimal_context_id(trace) if _uses_decimal(unwrap_deep(args)) else None
+    add_node(trace, op, args, kwargs, Raised.of(exc), GUARD, ctx=ctx, fn=fn)
+
+
+def project(trace: Trace, node: Node) -> Any:
+    """Return a copy of the container ``node.value`` whose items are traced."""
+    value = node.value
+    add_node(trace, "shape", (node,), None, shape(value), GUARD)
+    if type(value) is dict:
+        return {k: record_result(trace, "getitem", (node, k), None, v) for k, v in value.items()}
+    items = [record_result(trace, "getitem", (node, i), None, v) for i, v in enumerate(value)]
+    if type(value) is list:
+        return items
+    return rebuild_tuple(value, items)
+
+
+def encode(obj: Any, trace: Trace) -> tuple[Any, Any, bool]:
+    """Split a call argument into ``(reference, plain, linked)``.
+
+    ``reference`` is what gets stored in ``Node.args``: traced values from
+    ``trace`` become their :class:`Node`; lists, tuples and dicts are walked.
+    ``plain`` is the same structure with every traced value unwrapped.
+    ``linked`` says whether any part of ``obj`` belongs to ``trace``.
+    """
+    if isinstance(obj, Traced):
+        node = obj._rt_node
+        if node.trace is trace:
+            return node, node.value, True
+        return node.value, node.value, False
+    kind = type(obj)
+    if kind is list or isinstance(obj, tuple):
+        refs: list[Any] = []
+        plains: list[Any] = []
+        linked = changed = False
+        for item in obj:
+            ref, plain, item_linked = encode(item, trace)
+            refs.append(ref)
+            plains.append(plain)
+            linked = linked or item_linked
+            changed = changed or plain is not item
+        if not changed:
+            return obj, obj, False
+        if kind is list:
+            return refs, plains, linked
+        return rebuild_tuple(obj, refs), rebuild_tuple(obj, plains), linked
+    if kind is dict:
+        ref_d: dict[Any, Any] = {}
+        plain_d: dict[Any, Any] = {}
+        linked = changed = False
+        for k, item in obj.items():
+            ref, plain, item_linked = encode(item, trace)
+            ref_d[k] = ref
+            plain_d[k] = plain
+            linked = linked or item_linked
+            changed = changed or plain is not item
+        if not changed:
+            return obj, obj, False
+        return ref_d, plain_d, linked
+    return obj, obj, False
+
+
+def unwrap_deep(obj: Any) -> Any:
+    """Unwrap traced values inside lists, tuples and dicts."""
+    if isinstance(obj, Traced):
+        return obj._rt_node.value
+    kind = type(obj)
+    if kind is list:
+        return [unwrap_deep(item) for item in obj]
+    if isinstance(obj, tuple):
+        return rebuild_tuple(obj, [unwrap_deep(item) for item in obj])
+    if kind is dict:
+        return {k: unwrap_deep(v) for k, v in obj.items()}
+    return obj
+
+
+def call(
+    op: str,
+    fn: Callable[..., Any],
+    args: tuple[Any, ...],
+    kwargs: dict[str, Any] | None = None,
+    *,
+    meta: dict[str, Any] | None = None,
+) -> Any:
+    """Apply ``fn`` to the plain form of the arguments and record the call.
+
+    Records nothing (and returns the plain result) unless a trace is active
+    and at least one argument belongs to it.
+    """
+    trace = _active.get()
+    if trace is None or isinstance(trace, Observer):
+        if trace is not None:
+            trace.touch(args, kwargs or {})
+        return fn(*unwrap_deep(args), **unwrap_deep(kwargs or {}))
+    refs, plains, linked = encode(args, trace)
+    kw_refs: dict[str, Any] | None = None
+    kw_plains: dict[str, Any] = {}
+    if kwargs:
+        kw_refs, kw_plains, kw_linked = encode(kwargs, trace)
+        linked = linked or kw_linked
+    if not linked:
+        return fn(*plains, **kw_plains)
+    try:
+        result = fn(*plains, **kw_plains)
+    except Exception as exc:
+        record_raise(trace, op, tuple(refs), kw_refs, exc)
+        raise
+    ctx = (
+        decimal_context_id(trace)
+        if type(result) is Decimal or _uses_decimal(plains) or _uses_decimal(kw_plains.values())
+        else None
+    )
+    return record_result(trace, op, tuple(refs), kw_refs, result, meta, ctx)
+
+
+def binary(op: str, fn: Callable[[Any, Any], Any], left: Any, right: Any) -> Any:
+    """Fast path of :func:`call` for two scalar operands."""
+    if isinstance(left, Traced):
+        left_node: Node | None = left._rt_node
+        left = left_node.value  # type: ignore[union-attr]
+    else:
+        left_node = None
+    if isinstance(right, Traced):
+        right_node: Node | None = right._rt_node
+        right = right_node.value  # type: ignore[union-attr]
+    else:
+        right_node = None
+    trace = _active.get()
+    if trace is None:
+        return fn(left, right)
+    if isinstance(trace, Observer):
+        trace.touch_node(left_node)
+        trace.touch_node(right_node)
+        return fn(left, right)
+    left_ref, right_ref, linked = left, right, False
+    if left_node is not None and left_node.trace is trace:
+        left_ref, linked = left_node, True
+    if right_node is not None and right_node.trace is trace:
+        right_ref, linked = right_node, True
+    if not linked:
+        return fn(left, right)
+    try:
+        result = fn(left, right)
+    except Exception as exc:
+        record_raise(trace, op, (left_ref, right_ref), None, exc)
+        raise
+    ctx = (
+        decimal_context_id(trace)
+        if type(result) is Decimal or type(left) is Decimal or type(right) is Decimal
+        else None
+    )
+    return record_result(trace, op, (left_ref, right_ref), None, result, None, ctx)
+
+
+def observe_value(op: str, fn: Callable[..., Any], *operands: Any) -> Any:
+    """Like :func:`binary` but always records a guard and returns a plain result.
+
+    Used for comparisons, truth tests and conversions that must return a
+    builtin type (``__bool__``, ``__int__``, ``__hash__``...).
+    """
+    trace = _active.get()
+    plains = [
+        operand._rt_node.value if isinstance(operand, Traced) else operand for operand in operands
+    ]
+    if trace is None:
+        return fn(*plains)
+    if isinstance(trace, Observer):
+        trace.touch(*operands)
+        return fn(*plains)
+    refs: list[Any] = []
+    linked = False
+    for operand, plain in zip(operands, plains, strict=True):
+        if isinstance(operand, Traced) and operand._rt_node.trace is trace:
+            refs.append(operand._rt_node)
+            linked = True
+        else:
+            refs.append(plain)
+    if not linked:
+        return fn(*plains)
+    try:
+        result = fn(*plains)
+    except Exception as exc:
+        record_raise(trace, op, tuple(refs), None, exc)
+        raise
+    add_node(trace, op, tuple(refs), None, result, GUARD)
+    return result