redroot 0.2.0__py3-none-any.whl

redroot/ops.py

@@ -0,0 +1,332 @@
+"""The operation registry.
+
+Every node in a trace names the operation that produced it. The registry maps
+that name to the plain Python function that computes it. The *same* function
+is used when the value is first computed and whenever the graph re-evaluates
+the node, so re-evaluation reproduces traced execution exactly.
+
+Built-in operations (arithmetic, comparisons, conversions, ``str`` and
+``Decimal`` methods) are registered on import. User functions are registered
+with :func:`redroot.traced`.
+"""
+
+from __future__ import annotations
+
+import bisect
+import heapq
+import math
+import operator
+import statistics
+from collections.abc import Callable, Mapping
+from dataclasses import dataclass
+from decimal import Decimal
+from typing import Any, Literal
+
+Style = Literal["infix", "prefix", "call", "method", "getitem", "contains"]
+
+
+@dataclass(frozen=True, slots=True)
+class OpSpec:
+    """How to compute, and how to display, one kind of operation.
+
+    Attributes:
+        name: The name recorded on nodes.
+        fn: Computes the operation from plain argument values.
+        replayable: ``False`` for operations whose result cannot be recomputed
+            (e.g. LLM calls, web lookups). Nodes that depend on them become
+            stale instead of being re-evaluated when their inputs change.
+        symbol: The operator symbol (infix/prefix styles) or method name
+            (method style) used by :meth:`redroot.Trace.explain`.
+        style: How :meth:`redroot.Trace.explain` renders the operation.
+    """
+
+    name: str
+    fn: Callable[..., Any]
+    replayable: bool = True
+    symbol: str | None = None
+    style: Style = "call"
+
+
+_REGISTRY: dict[str, OpSpec] = {}
+
+
+def register_op(
+    name: str,
+    fn: Callable[..., Any],
+    *,
+    replayable: bool = True,
+    symbol: str | None = None,
+    style: Style = "call",
+) -> OpSpec:
+    """Register (or replace) the operation ``name``.
+
+    Operations must be pure: given equal arguments they must return equal
+    results and have no side effects, otherwise re-evaluation is meaningless.
+    Mark anything else ``replayable=False``.
+    """
+    spec = OpSpec(name, fn, replayable, symbol, style)
+    _REGISTRY[name] = spec
+    return spec
+
+
+def get_op(name: str) -> OpSpec | None:
+    """Return the registered operation ``name``, or ``None``."""
+    return _REGISTRY.get(name)
+
+
+def registered_ops() -> Mapping[str, OpSpec]:
+    """Return a read-only snapshot of the registry."""
+    return dict(_REGISTRY)
+
+
+class Opaque:
+    """Stands in for a constant that could not be serialized.
+
+    Evaluating a node with an opaque argument fails, which marks the node
+    stale rather than silently producing a wrong value.
+    """
+
+    __slots__ = ("description",)
+
+    def __init__(self, description: str) -> None:
+        self.description = description
+
+    def __repr__(self) -> str:
+        return f"<opaque {self.description}>"
+
+
+def shape(value: Any) -> Any:
+    """The structural signature of a container: its keys or its length."""
+    return tuple(value) if isinstance(value, dict) else len(value)
+
+
+def observe(value: Any) -> Any:
+    """Identity; records that a value was consumed in a way we cannot model."""
+    return value
+
+
+def fstring(*parts: Any) -> str:
+    """Rebuild an f-string from literal parts and ``(value, conversion, spec)`` triples."""
+    out: list[str] = []
+    for part in parts:
+        if isinstance(part, str):
+            out.append(part)
+            continue
+        value, conversion, spec = part
+        if conversion == "r":
+            value = repr(value)
+        elif conversion == "s":
+            value = str(value)
+        elif conversion == "a":
+            value = ascii(value)
+        out.append(format(value, spec))
+    return "".join(out)
+
+
+BINARY_OPS: dict[str, tuple[Callable[[Any, Any], Any], str]] = {
+    "add": (operator.add, "+"),
+    "sub": (operator.sub, "-"),
+    "mul": (operator.mul, "*"),
+    "truediv": (operator.truediv, "/"),
+    "floordiv": (operator.floordiv, "//"),
+    "mod": (operator.mod, "%"),
+    "pow": (pow, "**"),
+    "lshift": (operator.lshift, "<<"),
+    "rshift": (operator.rshift, ">>"),
+    "and": (operator.and_, "&"),
+    "or": (operator.or_, "|"),
+    "xor": (operator.xor, "^"),
+}
+
+COMPARISON_OPS: dict[str, tuple[Callable[[Any, Any], Any], str]] = {
+    "lt": (operator.lt, "<"),
+    "le": (operator.le, "<="),
+    "gt": (operator.gt, ">"),
+    "ge": (operator.ge, ">="),
+    "eq": (operator.eq, "=="),
+    "ne": (operator.ne, "!="),
+}
+
+UNARY_OPS: dict[str, tuple[Callable[[Any], Any], str]] = {
+    "neg": (operator.neg, "-"),
+    "pos": (operator.pos, "+"),
+    "invert": (operator.invert, "~"),
+}
+
+# Builtins applied as functions: name -> implementation.
+FUNCTION_OPS: dict[str, Callable[..., Any]] = {
+    "abs": abs,
+    "round": round,
+    "divmod": divmod,
+    "floor": math.floor,
+    "ceil": math.ceil,
+    "trunc": math.trunc,
+    "bool": bool,
+    "int": int,
+    "float": float,
+    "complex": complex,
+    "str": str,
+    "decimal": Decimal,
+    "format": format,
+    "hash": hash,
+    "len": len,
+    "index": operator.index,
+    "shape": shape,
+    "observe": observe,
+    "fstring": fstring,
+    "min": min,
+    "max": max,
+    "sum": sum,
+    "range": range,
+    "repr": repr,
+    "ascii": ascii,
+    "sorted": sorted,
+    "bisect.bisect_left": bisect.bisect_left,
+    "bisect.bisect_right": bisect.bisect_right,
+    "heapq.nsmallest": heapq.nsmallest,
+    "heapq.nlargest": heapq.nlargest,
+}
+
+STATISTICS_FUNCTIONS: tuple[str, ...] = (
+    "fmean",
+    "geometric_mean",
+    "harmonic_mean",
+    "mean",
+    "median",
+    "median_high",
+    "median_low",
+    "mode",
+    "pstdev",
+    "pvariance",
+    "stdev",
+    "variance",
+)
+FUNCTION_OPS.update(
+    {f"statistics.{name}": getattr(statistics, name) for name in STATISTICS_FUNCTIONS}
+)
+
+MATH_FUNCTIONS: tuple[str, ...] = (
+    "copysign",
+    "exp",
+    "fabs",
+    "fsum",
+    "hypot",
+    "isclose",
+    "isfinite",
+    "isinf",
+    "isnan",
+    "log",
+    "log10",
+    "log2",
+    "pow",
+    "prod",
+    "sqrt",
+)
+FUNCTION_OPS.update({f"math.{name}": getattr(math, name) for name in MATH_FUNCTIONS})
+
+STR_METHODS: tuple[str, ...] = (
+    # str -> str
+    "capitalize",
+    "casefold",
+    "center",
+    "expandtabs",
+    "format",
+    "format_map",
+    "join",
+    "ljust",
+    "lower",
+    "lstrip",
+    "removeprefix",
+    "removesuffix",
+    "replace",
+    "rjust",
+    "rstrip",
+    "strip",
+    "swapcase",
+    "title",
+    "translate",
+    "upper",
+    "zfill",
+    # str -> containers / ints / bools
+    "count",
+    "endswith",
+    "find",
+    "index",
+    "isalnum",
+    "isalpha",
+    "isascii",
+    "isdecimal",
+    "isdigit",
+    "isidentifier",
+    "islower",
+    "isnumeric",
+    "isprintable",
+    "isspace",
+    "istitle",
+    "isupper",
+    "partition",
+    "rfind",
+    "rindex",
+    "rpartition",
+    "rsplit",
+    "split",
+    "splitlines",
+    "startswith",
+    "encode",
+)
+
+DECIMAL_METHODS: tuple[str, ...] = (
+    "adjusted",
+    "as_integer_ratio",
+    "as_tuple",
+    "compare",
+    "copy_abs",
+    "copy_negate",
+    "copy_sign",
+    "exp",
+    "fma",
+    "is_finite",
+    "is_infinite",
+    "is_nan",
+    "is_signed",
+    "is_zero",
+    "ln",
+    "log10",
+    "max",
+    "min",
+    "normalize",
+    "quantize",
+    "scaleb",
+    "sqrt",
+    "to_integral",
+    "to_integral_exact",
+    "to_integral_value",
+)
+
+FLOAT_METHODS: tuple[str, ...] = ("as_integer_ratio", "is_integer")
+INT_METHODS: tuple[str, ...] = ("as_integer_ratio", "bit_count", "bit_length")
+
+
+def _register_builtins() -> None:
+    for name, (fn, symbol) in BINARY_OPS.items():
+        register_op(name, fn, symbol=symbol, style="infix")
+    for name, (fn, symbol) in COMPARISON_OPS.items():
+        register_op(name, fn, symbol=symbol, style="infix")
+    for name, (unary, symbol) in UNARY_OPS.items():
+        register_op(name, unary, symbol=symbol, style="prefix")
+    for name, fn in FUNCTION_OPS.items():
+        register_op(name, fn)
+    register_op("decimal", Decimal, symbol="Decimal")
+    register_op("getitem", operator.getitem, style="getitem")
+    register_op("contains", operator.contains, symbol="in", style="contains")
+    for prefix, owner, methods in (
+        ("str", str, STR_METHODS),
+        ("decimal", Decimal, DECIMAL_METHODS),
+        ("float", float, FLOAT_METHODS),
+        ("int", int, INT_METHODS),
+    ):
+        for method in methods:
+            register_op(f"{prefix}.{method}", getattr(owner, method), symbol=method, style="method")
+
+
+_register_builtins()

redroot/paths.py

@@ -0,0 +1,162 @@
+"""Semantic keys for values inside nested data.
+
+A key names one value by its position in a nested structure of dicts and
+lists, prefixed by a *root* that says which structure it belongs to::
+
+    ext:bank.line_items[3].amount
+    out:form1["line 12"]
+    out:tables[0][2]
+
+Dict keys that are ASCII identifiers use dot notation. Other keys and list indexes use
+brackets holding a JSON literal, so ``[3]`` is the integer 3 (a list index or
+an ``int`` dict key) and ``["3"]`` is the string ``"3"``. Keys round-trip:
+``parse_key(format_key(root, path)) == (root, path)``.
+"""
+
+from __future__ import annotations
+
+import copy
+import json
+import operator
+import re
+from collections.abc import Mapping, Sequence
+from typing import Any
+
+from redroot._core import rebuild_tuple, unwrap
+
+__all__ = ["apply_edits", "child_key", "format_key", "get_path", "parse_key", "set_path"]
+
+Segment = str | int | float | bool | None
+Path = tuple[Segment, ...]
+
+_ROOT = re.compile(r"[A-Za-z_][A-Za-z0-9_\-]*")
+_NAME = re.compile(r"[A-Za-z_][A-Za-z0-9_]*")
+
+
+def _check_root(root: str) -> str:
+    if not _ROOT.fullmatch(root):
+        raise ValueError(f"invalid root {root!r}: expected a name such as 'ext' or 'out'")
+    return root
+
+
+def _segment(segment: Any, first: bool) -> str:
+    segment = unwrap(segment)
+    if isinstance(segment, str) and _NAME.fullmatch(segment):
+        return segment if first else f".{segment}"
+    if segment is None or isinstance(segment, (str, int, float, bool)):
+        return f"[{json.dumps(segment)}]"
+    # Keys that are not JSON scalars (tuples, enums...) are named by their repr.
+    return f"[{json.dumps(repr(segment))}]"
+
+
+def child_key(parent: str, segment: Any) -> str:
+    """Return the key of ``segment`` inside the value named ``parent``.
+
+    >>> child_key("ext", "bank"), child_key("ext:bank", 0)
+    ('ext:bank', 'ext:bank[0]')
+    """
+    if ":" not in parent:  # a bare root
+        return f"{parent}:{_segment(segment, first=True)}"
+    return parent + _segment(segment, first=False)
+
+
+def format_key(root: str, path: Sequence[Any] = ()) -> str:
+    """Build the key of the value at ``path`` under ``root``.
+
+    >>> format_key("ext", ["bank", "line_items", 3, "amount"])
+    'ext:bank.line_items[3].amount'
+    """
+    _check_root(root)
+    if not path:
+        return root
+    return f"{root}:" + "".join(_segment(s, first=i == 0) for i, s in enumerate(path))
+
+
+_DECODER = json.JSONDecoder()
+
+
+def parse_key(key: str) -> tuple[str, Path]:
+    """Split a key into its root and path.
+
+    >>> parse_key('out:form1["line 12"]')
+    ('out', ('form1', 'line 12'))
+    """
+    root, sep, rest = key.partition(":")
+    _check_root(root)
+    if not sep:
+        return root, ()
+    path: list[Segment] = []
+    i = 0
+    while i < len(rest):
+        if rest[i] == "[":
+            try:
+                value, end = _DECODER.raw_decode(rest, i + 1)
+            except json.JSONDecodeError as exc:
+                raise ValueError(f"invalid key {key!r}: bad bracket segment") from exc
+            if end >= len(rest) or rest[end] != "]":
+                raise ValueError(f"invalid key {key!r}: unclosed bracket")
+            if isinstance(value, (list, dict)):
+                raise ValueError(f"invalid key {key!r}: brackets must hold a scalar")
+            path.append(value)
+            i = end + 1
+            continue
+        if rest[i] == ".":
+            if not path:
+                raise ValueError(f"invalid key {key!r}: unexpected '.'")
+            i += 1
+        elif path:
+            raise ValueError(f"invalid key {key!r}: expected '.' or '['")
+        match = _NAME.match(rest, i)
+        if match is None:
+            raise ValueError(f"invalid key {key!r}: expected a name at position {i}")
+        path.append(match.group())
+        i = match.end()
+    if not path:
+        raise ValueError(f"invalid key {key!r}: empty path")
+    return root, tuple(path)
+
+
+def get_path(data: Any, path: Sequence[Segment]) -> Any:
+    """Return the value at ``path`` inside ``data``."""
+    for segment in path:
+        data = data[segment]
+    return data
+
+
+def set_path(data: Any, path: Sequence[Segment], value: Any) -> None:
+    """Replace the value at ``path`` inside ``data`` (in place)."""
+    if not path:
+        raise ValueError("cannot replace the root value in place")
+    parent = get_path(data, path[:-1])
+    if isinstance(parent, tuple):
+        raise TypeError(f"cannot assign into a tuple at {list(path[:-1])}")
+    parent[path[-1]] = value
+
+
+def apply_edits(data: Any, edits: Mapping[str, Any], *, root: str) -> Any:
+    """Return a deep copy of ``data`` with each edit applied.
+
+    Edits are keyed like :meth:`redroot.Trace.propagate` edits. This is how
+    to build the inputs for re-executing a workflow after a review.
+    """
+    edited = copy.deepcopy(data)
+    for key, value in edits.items():
+        key_root, path = parse_key(key)
+        if key_root != root:
+            raise ValueError(f"edit {key!r} does not belong to root {root!r}")
+        edited = _replace(edited, path, value)
+    return edited
+
+
+def _replace(data: Any, path: Sequence[Segment], value: Any) -> Any:
+    """Return ``data`` with the value at ``path`` replaced, rebuilding tuples."""
+    if not path:
+        return value
+    head = path[0]
+    child = _replace(data[head], path[1:], value)
+    if isinstance(data, tuple):
+        items = list(data)
+        items[operator.index(head)] = child  # type: ignore[arg-type]
+        return rebuild_tuple(data, items)
+    data[head] = child
+    return data