redroot 0.2.0__py3-none-any.whl

redroot/instrument.py

@@ -0,0 +1,614 @@
+"""Source instrumentation: trace what operator overloading cannot see.
+
+Traced values are subclasses of builtins, so most operations reach their
+overridden methods. A few do not, because CPython handles them in C without
+consulting the subclass:
+
+* a plain ``float``/``Decimal`` on the left of a traced ``int``
+  (``0.8 * income``) and the matching comparisons (``0.5 < size``);
+* conversions that parse or read the raw value: ``float(text)``,
+  ``Decimal(text)``, ``int(text)``, ``math.sqrt(x)``;
+* methods of *plain* strings given traced arguments: ``", ".join(parts)``,
+  ``"{:,.2f}".format(x)``, and f-strings;
+* indexing a plain list with a traced ``int`` (``table[size]``) and
+  ``range(n)``.
+
+Values computed that way silently lose their lineage. For code you execute
+from source (such as generated workflow code), :func:`exec_source` rewrites
+the syntax tree so these operations are routed through RedRoot first, the
+way pytest rewrites ``assert`` statements. Instrumented code behaves exactly
+like the original; with no active trace it only adds a function call per
+operation.
+
+Usage::
+
+    namespace = redroot.instrument.exec_source(generated_code)
+    trace, result = redroot.run(namespace["main"], extracted, input_root="ext")
+"""
+
+from __future__ import annotations
+
+import ast
+import bisect
+import collections
+import heapq
+import itertools
+import math
+import operator
+import statistics
+from collections.abc import Callable
+from decimal import Decimal
+from types import CodeType
+from typing import Any
+
+from redroot import _core
+from redroot._core import Traced, observe_value
+from redroot.ops import (
+    BINARY_OPS,
+    COMPARISON_OPS,
+    DECIMAL_METHODS,
+    MATH_FUNCTIONS,
+    STATISTICS_FUNCTIONS,
+    STR_METHODS,
+    fstring,
+)
+
+__all__ = ["RUNTIME_NAME", "compile_source", "exec_source", "instrument"]
+
+RUNTIME_NAME = "__redroot__"
+"""Global name under which instrumented code finds the runtime helpers."""
+
+# --------------------------------------------------------------------------
+# Runtime helpers called by instrumented code
+# --------------------------------------------------------------------------
+
+# Pure callables recorded as one node when given traced arguments: id -> op.
+_CALLS: dict[int, str] = {
+    id(fn): name
+    for fn, name in [
+        (int, "int"),
+        (float, "float"),
+        (str, "str"),
+        (bool, "bool"),
+        (Decimal, "decimal"),
+        (round, "round"),
+        (abs, "abs"),
+        (min, "min"),
+        (max, "max"),
+        (sum, "sum"),
+        (len, "len"),
+        (divmod, "divmod"),
+        (pow, "pow"),
+        (format, "format"),
+        (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"),
+        (math.floor, "floor"),
+        (math.ceil, "ceil"),
+        (math.trunc, "trunc"),
+        *((getattr(math, name), f"math.{name}") for name in MATH_FUNCTIONS),
+        *((getattr(statistics, name), f"statistics.{name}") for name in STATISTICS_FUNCTIONS),
+    ]
+}
+# Position of the iterable argument to materialize so it can be inspected
+# (and is not consumed twice); min/max only take one when called with one.
+_ITERABLE_ARG: dict[int, int] = {
+    id(sum): 0,
+    id(sorted): 0,
+    id(math.fsum): 0,
+    id(math.prod): 0,
+    id(heapq.nsmallest): 1,
+    id(heapq.nlargest): 1,
+    **{id(getattr(statistics, name)): 0 for name in STATISTICS_FUNCTIONS},
+}
+_SINGLE_ITERABLE = {id(min), id(max)}
+_METHODS: dict[type, tuple[str, frozenset[str]]] = {
+    str: ("str", frozenset(STR_METHODS)),
+    Decimal: ("decimal", frozenset(DECIMAL_METHODS)),
+}
+# Containers whose methods only store, fetch or compare values through the
+# values' own (traced) methods.
+_CONTAINERS: tuple[type, ...] = (
+    list,
+    dict,
+    set,
+    frozenset,
+    tuple,
+    collections.deque,
+    collections.defaultdict,
+    collections.OrderedDict,
+    collections.Counter,
+)
+# Builtins that pass values through, or depend only on types, so a call that
+# returns no traced value has not lost lineage.
+_TRANSPARENT: set[int] = {
+    id(fn)
+    for fn in (
+        print,
+        isinstance,
+        issubclass,
+        callable,
+        hasattr,
+        getattr,
+        setattr,
+        delattr,
+        id,
+        iter,
+        next,
+        enumerate,
+        zip,
+        reversed,
+        map,
+        filter,
+        list,
+        tuple,
+        dict,
+        set,
+        frozenset,
+        any,
+        all,
+        vars,
+        dir,
+        object,
+        type,
+        slice,
+    )
+}
+_TRANSPARENT_MODULES = frozenset(
+    {"logging", "warnings", "copy", "itertools", "collections", "dataclasses"}
+)
+_INSTRUMENTED_FILES: set[str] = set()
+
+
+def _has_traced(obj: Any, depth: int = 2) -> bool:
+    if isinstance(obj, Traced):
+        return True
+    if depth and type(obj) in (list, tuple):
+        return any(_has_traced(item, depth - 1) for item in obj)
+    if depth and type(obj) is dict:
+        return any(_has_traced(item, depth - 1) for item in obj.values())
+    return False
+
+
+def _pure_callable(func: Any) -> str | None:
+    """The op name of a callable RedRoot can record as one node, if any."""
+    op = _CALLS.get(id(func))
+    if op is not None:
+        return op
+    owner = getattr(func, "__objclass__", None)  # unbound method, e.g. str.lower
+    method = _METHODS.get(owner) if isinstance(owner, type) else None
+    if method is not None and getattr(func, "__name__", None) in method[1]:
+        return f"{method[0]}.{func.__name__}"
+    return None
+
+
+def _wrap_callable(func: Any) -> Any:
+    """Route calls of a known pure callable through :func:`call`."""
+    if func is None or _pure_callable(func) is None:
+        return func
+
+    def routed(*args: Any, **kwargs: Any) -> Any:
+        return call(func, *args, **kwargs)
+
+    return routed
+
+
+def _is_instrumented(func: Any) -> bool:
+    code = getattr(func, "__code__", None) or getattr(
+        getattr(func, "__func__", None), "__code__", None
+    )
+    return code is not None and code.co_filename in _INSTRUMENTED_FILES
+
+
+def _is_transparent(func: Any) -> bool:
+    if id(func) in _TRANSPARENT or hasattr(func, "__redroot_op__") or _is_instrumented(func):
+        return True
+    if isinstance(getattr(func, "__self__", None), _CONTAINERS):
+        return True
+    module = getattr(func, "__module__", None) or ""
+    return module.partition(".")[0] in _TRANSPARENT_MODULES
+
+
+def _keeps_lineage(result: Any, trace: Any, first_new_node: int, depth: int = 2) -> bool:
+    """Whether ``result`` carries traced values that link back to the trace's inputs."""
+    if isinstance(result, Traced):
+        node = result._rt_node
+        # A value rebuilt from a plain number (e.g. type(x)(...) in a library)
+        # becomes a fresh anonymous input: that is lost lineage, not kept.
+        fresh_input = node.kind == "leaf" and node.key is None and node.id >= first_new_node
+        return node.trace is trace and not fresh_input
+    if not depth:
+        return False
+    if type(result) in (list, tuple, set, frozenset) or isinstance(result, tuple):
+        return any(_keeps_lineage(item, trace, first_new_node, depth - 1) for item in result)
+    if type(result) is dict:
+        return any(
+            _keeps_lineage(item, trace, first_new_node, depth - 1) for item in result.values()
+        )
+    attributes = getattr(result, "__dict__", None)
+    if isinstance(attributes, dict):
+        return any(
+            _keeps_lineage(item, trace, first_new_node, depth - 1) for item in attributes.values()
+        )
+    return False
+
+
+def _sort_in_place(items: list[Any], key: Any = None, reverse: bool = False) -> None:
+    if key is None and _has_traced(items, depth=1):
+        items[:] = _core.call(
+            "sorted", sorted, (list(items),), {"reverse": reverse} if reverse else None
+        )
+    else:
+        items.sort(key=_wrap_callable(key), reverse=reverse)
+
+
+def _opaque_call(func: Any, args: tuple[Any, ...], kwargs: dict[str, Any]) -> Any:
+    """Call code RedRoot cannot see into; guard its traced inputs if lineage is lost.
+
+    The guard cannot be re-evaluated, so any change to those inputs means
+    re-execution: sound, if conservative.
+    """
+    trace = _core.active_trace()
+    if trace is None or _is_transparent(func):
+        return func(*args, **kwargs)
+    refs, _, linked = _core.encode(args, trace)
+    kw_refs, _, kw_linked = _core.encode(kwargs, trace)
+    if not (linked or kw_linked):
+        return func(*args, **kwargs)
+    name = getattr(func, "__qualname__", None) or type(func).__qualname__
+    op = f"call:{getattr(func, '__module__', None) or ''}.{name}".replace(":.", ":")
+    first_new_node = len(trace.nodes)
+    try:
+        result = func(*args, **kwargs)
+    except Exception as exc:
+        _core.record_raise(trace, op, tuple(refs), kw_refs or None, exc, _core.NOT_REPLAYABLE)
+        raise
+    if not _keeps_lineage(result, trace, first_new_node):
+        _core.add_node(
+            trace,
+            op,
+            tuple(refs),
+            kw_refs or None,
+            _core.unwrap_deep(result),
+            _core.GUARD,
+            fn=_core.NOT_REPLAYABLE,
+        )
+    return result
+
+
+def call(func: Any, /, *args: Any, **kwargs: Any) -> Any:
+    """Call ``func`` from instrumented code, recording what overloading cannot see."""
+    func_id = id(func)
+    op = _CALLS.get(func_id)
+    if op is not None:
+        if kwargs.get("key") is not None:
+            # A custom key may close over traced values: run natively so the
+            # comparisons of its results are recorded as guards.
+            kwargs["key"] = _wrap_callable(kwargs["key"])
+            return func(*args, **kwargs)
+        position = _ITERABLE_ARG.get(
+            func_id, 0 if func_id in _SINGLE_ITERABLE and len(args) == 1 else -1
+        )
+        if 0 <= position < len(args) and type(args[position]) not in (list, tuple):
+            args = (*args[:position], list(args[position]), *args[position + 1 :])
+        if _has_traced(args) or _has_traced(kwargs):
+            return _core.call(op, func, args, kwargs or None)
+        return func(*args, **kwargs)
+    if func is map or func is filter:
+        if args:
+            args = (_wrap_callable(args[0]), *args[1:])
+        return func(*args, **kwargs)
+    if func is type and len(args) == 1 and not kwargs and isinstance(args[0], Traced):
+        return args[0]._base  # what type() returns for the plain value
+    receiver = getattr(func, "__self__", None)
+    method = _METHODS.get(type(receiver))
+    if method is not None and func.__name__ in method[1]:
+        if func.__name__ == "join" and args and type(args[0]) not in (list, tuple):
+            args = (list(args[0]),)
+        if _has_traced(args) or _has_traced(kwargs):
+            name = func.__name__
+            unbound = getattr(type(receiver), name)
+            return _core.call(f"{method[0]}.{name}", unbound, (receiver, *args), kwargs or None)
+        return func(*args, **kwargs)
+    if type(receiver) is list and getattr(func, "__name__", None) == "sort":
+        return _sort_in_place(receiver, *args, **kwargs)
+    op = _pure_callable(func)  # unbound methods: str.upper(x)
+    if op is not None and (_has_traced(args) or _has_traced(kwargs)):
+        return _core.call(op, func, args, kwargs or None)
+    if not (_has_traced(args) or _has_traced(kwargs)):
+        return func(*args, **kwargs)
+    return _opaque_call(func, args, kwargs)
+
+
+def binop(op: str, left: Any, right: Any) -> Any:
+    """``left <op> right``, recorded if either operand is traced."""
+    fn = BINARY_OPS[op][0]
+    if isinstance(left, Traced) or isinstance(right, Traced):
+        return _core.binary(op, fn, left, right)
+    if op == "mod" and type(left) is str and _has_traced(right):
+        return _core.call(op, fn, (left, right))  # "%s" % (traced, ...)
+    return fn(left, right)
+
+
+_INPLACE: dict[str, Callable[[Any, Any], Any]] = {
+    "add": operator.iadd,
+    "sub": operator.isub,
+    "mul": operator.imul,
+    "truediv": operator.itruediv,
+    "floordiv": operator.ifloordiv,
+    "mod": operator.imod,
+    "pow": operator.ipow,
+    "lshift": operator.ilshift,
+    "rshift": operator.irshift,
+    "and": operator.iand,
+    "or": operator.ior,
+    "xor": operator.ixor,
+}
+
+
+def ibinop(op: str, left: Any, right: Any) -> Any:
+    """``left <op>= right``: in place for mutable ``left``, recorded for traced values."""
+    if isinstance(left, Traced) or (
+        isinstance(right, Traced) and not hasattr(type(left), f"__i{op}__")
+    ):
+        return _core.binary(op, BINARY_OPS[op][0], left, right)
+    return _INPLACE[op](left, right)
+
+
+def compare(op: str, left: Any, right: Any) -> Any:
+    """``left <op> right`` for a single comparison, recorded as a guard."""
+    if op in ("in", "not in"):
+        if isinstance(left, Traced) and not isinstance(right, Traced):
+            left = observe_value("observe", _identity, left)  # membership in a plain container
+        result = left in right
+        return not result if op == "not in" else result
+    if isinstance(left, Traced) or isinstance(right, Traced):
+        return observe_value(op, COMPARISON_OPS[op][0], left, right)
+    return COMPARISON_OPS[op][0](left, right)
+
+
+def _identity(value: Any) -> Any:
+    return value
+
+
+def _plain_key(key: Any) -> Any:
+    if isinstance(key, Traced):
+        return observe_value("observe", _identity, key)
+    if type(key) is slice and _has_traced((key.start, key.stop, key.step)):
+        return slice(_plain_key(key.start), _plain_key(key.stop), _plain_key(key.step))
+    if type(key) is tuple and _has_traced(key, depth=1):
+        return tuple(_plain_key(k) for k in key)
+    return key
+
+
+def key(value: Any) -> Any:
+    """A subscript key for a store or delete: traced keys become guarded plain keys."""
+    return _plain_key(value)
+
+
+def getitem(container: Any, index: Any) -> Any:
+    """``container[index]``; a traced index into a plain container is guarded."""
+    if not isinstance(container, Traced):
+        index = _plain_key(index)
+    return container[index]
+
+
+def fstring_(*parts: Any) -> Any:
+    """Build an f-string, recording it if any interpolated value is traced."""
+    if _has_traced(parts, depth=8):  # values may be containers of traced values
+        return _core.call("fstring", fstring, parts)
+    return fstring(*parts)
+
+
+class _Runtime:
+    """Namespace object bound to :data:`RUNTIME_NAME` in instrumented code."""
+
+    call = staticmethod(call)
+    binop = staticmethod(binop)
+    ibinop = staticmethod(ibinop)
+    compare = staticmethod(compare)
+    getitem = staticmethod(getitem)
+    key = staticmethod(key)
+    fstring = staticmethod(fstring_)
+    slice = slice
+
+
+RUNTIME = _Runtime()
+
+# --------------------------------------------------------------------------
+# AST transformation
+# --------------------------------------------------------------------------
+
+_BINOP_NAMES: dict[type[ast.operator], str] = {
+    ast.Add: "add",
+    ast.Sub: "sub",
+    ast.Mult: "mul",
+    ast.Div: "truediv",
+    ast.FloorDiv: "floordiv",
+    ast.Mod: "mod",
+    ast.Pow: "pow",
+    ast.LShift: "lshift",
+    ast.RShift: "rshift",
+    ast.BitAnd: "and",
+    ast.BitOr: "or",
+    ast.BitXor: "xor",
+}
+_COMPARE_NAMES: dict[type[ast.cmpop], str] = {
+    ast.Lt: "lt",
+    ast.LtE: "le",
+    ast.Gt: "gt",
+    ast.GtE: "ge",
+    ast.Eq: "eq",
+    ast.NotEq: "ne",
+    ast.In: "in",
+    ast.NotIn: "not in",
+}
+# Calls that depend on the calling frame must not be wrapped.
+_FRAME_SENSITIVE = frozenset(
+    {"super", "locals", "vars", "globals", "eval", "exec", "dir", "breakpoint"}
+)
+_CONVERSIONS = {-1: None, 115: "s", 114: "r", 97: "a"}
+
+
+def _runtime(attr: str) -> ast.expr:
+    return ast.Attribute(ast.Name(RUNTIME_NAME, ast.Load()), attr, ast.Load())
+
+
+def _rt_call(attr: str, *args: ast.expr) -> ast.Call:
+    return ast.Call(_runtime(attr), list(args), [])
+
+
+def _is_constant(node: ast.expr) -> bool:
+    return isinstance(node, ast.Constant)
+
+
+class _Instrumenter(ast.NodeTransformer):
+    def __init__(self) -> None:
+        self._temps = itertools.count()
+
+    def _temp(self) -> str:
+        return f"__redroot_tmp{next(self._temps)}"
+
+    # Annotations are left alone: rewriting them gains nothing.
+    def visit_arg(self, node: ast.arg) -> ast.arg:
+        return node
+
+    def visit_AnnAssign(self, node: ast.AnnAssign) -> ast.AnnAssign:
+        node.target = self.visit(node.target)
+        if node.value is not None:
+            node.value = self.visit(node.value)
+        return node
+
+    def _visit_function(self, node: ast.FunctionDef | ast.AsyncFunctionDef) -> ast.AST:
+        returns, node.returns = node.returns, None
+        self.generic_visit(node)
+        node.returns = returns
+        return node
+
+    visit_FunctionDef = _visit_function
+    visit_AsyncFunctionDef = _visit_function
+
+    def visit_BinOp(self, node: ast.BinOp) -> ast.expr:
+        self.generic_visit(node)
+        name = _BINOP_NAMES.get(type(node.op))
+        if name is None or (_is_constant(node.left) and _is_constant(node.right)):
+            return node
+        return ast.copy_location(_rt_call("binop", ast.Constant(name), node.left, node.right), node)
+
+    def visit_Compare(self, node: ast.Compare) -> ast.expr:
+        self.generic_visit(node)
+        if len(node.ops) != 1:
+            return node  # chained comparisons keep their overloading-based tracing
+        name = _COMPARE_NAMES.get(type(node.ops[0]))
+        if name is None or (_is_constant(node.left) and _is_constant(node.comparators[0])):
+            return node
+        call = _rt_call("compare", ast.Constant(name), node.left, node.comparators[0])
+        return ast.copy_location(call, node)
+
+    def visit_Call(self, node: ast.Call) -> ast.expr:
+        self.generic_visit(node)
+        if isinstance(node.func, ast.Name) and node.func.id in _FRAME_SENSITIVE:
+            return node
+        call = ast.Call(_runtime("call"), [node.func, *node.args], node.keywords)
+        return ast.copy_location(call, node)
+
+    def visit_JoinedStr(self, node: ast.JoinedStr) -> ast.expr:
+        parts: list[ast.expr] = []
+        for value in node.values:
+            if isinstance(value, ast.FormattedValue):
+                spec: ast.expr = (
+                    self.visit(value.format_spec)
+                    if value.format_spec is not None
+                    else ast.Constant("")
+                )
+                conversion = ast.Constant(_CONVERSIONS.get(value.conversion))
+                parts.append(ast.Tuple([self.visit(value.value), conversion, spec], ast.Load()))
+            else:
+                parts.append(value)
+        return ast.copy_location(_rt_call("fstring", *parts), node)
+
+    def _key(self, index: ast.expr) -> ast.expr:
+        if isinstance(index, ast.Slice):
+            bounds = [
+                b if b is not None else ast.Constant(None)
+                for b in (index.lower, index.upper, index.step)
+            ]
+            index = ast.Call(_runtime("slice"), bounds, [])
+        return index
+
+    def visit_Subscript(self, node: ast.Subscript) -> ast.expr:
+        self.generic_visit(node)
+        index = self._key(node.slice)
+        if isinstance(node.ctx, ast.Load):
+            return ast.copy_location(_rt_call("getitem", node.value, index), node)
+        node.slice = _rt_call("key", index)
+        return node
+
+    def visit_AugAssign(self, node: ast.AugAssign) -> ast.AST | list[ast.stmt]:
+        name = _BINOP_NAMES.get(type(node.op))
+        if name is None:
+            return self.generic_visit(node)
+        value = self.visit(node.value)
+        target = node.target
+        prelude: list[ast.stmt] = []
+        if isinstance(target, ast.Name):
+            load: ast.expr = ast.Name(target.id, ast.Load())
+            store: ast.expr = ast.Name(target.id, ast.Store())
+        elif isinstance(target, ast.Attribute):
+            obj = self._temp()
+            prelude.append(ast.Assign([ast.Name(obj, ast.Store())], self.visit(target.value)))
+            load = ast.Attribute(ast.Name(obj, ast.Load()), target.attr, ast.Load())
+            store = ast.Attribute(ast.Name(obj, ast.Load()), target.attr, ast.Store())
+        elif isinstance(target, ast.Subscript):
+            obj, idx = self._temp(), self._temp()
+            prelude.append(ast.Assign([ast.Name(obj, ast.Store())], self.visit(target.value)))
+            index = _rt_call("key", self._key(self.visit(target.slice)))
+            prelude.append(ast.Assign([ast.Name(idx, ast.Store())], index))
+            load = ast.Subscript(ast.Name(obj, ast.Load()), ast.Name(idx, ast.Load()), ast.Load())
+            store = ast.Subscript(ast.Name(obj, ast.Load()), ast.Name(idx, ast.Load()), ast.Store())
+        else:  # pragma: no cover - the grammar allows no other targets
+            return self.generic_visit(node)
+        assign = ast.Assign([store], _rt_call("ibinop", ast.Constant(name), load, value))
+        statements = [*prelude, assign]
+        for statement in statements:
+            ast.copy_location(statement, node)
+        return statements
+
+
+def instrument(source: str, filename: str = "<workflow>") -> ast.Module:
+    """Parse ``source`` and return its instrumented syntax tree."""
+    tree: ast.Module = _Instrumenter().visit(ast.parse(source, filename=filename))
+    return ast.fix_missing_locations(tree)
+
+
+def compile_source(source: str, filename: str = "<workflow>") -> CodeType:
+    """Compile ``source`` with instrumentation, for use with ``exec``.
+
+    The code expects :data:`RUNTIME_NAME` in its globals; prefer
+    :func:`exec_source`, which provides it.
+    """
+    _INSTRUMENTED_FILES.add(filename)
+    # dont_inherit: never leak this module's __future__ flags into the workflow.
+    return compile(instrument(source, filename), filename, "exec", dont_inherit=True)
+
+
+def exec_source(
+    source: str,
+    namespace: dict[str, Any] | None = None,
+    *,
+    filename: str = "<workflow>",
+) -> dict[str, Any]:
+    """Execute instrumented ``source`` and return its namespace.
+
+    Functions defined by the source stay instrumented whenever they are
+    called later, e.g. by :func:`redroot.run`.
+    """
+    namespace = {} if namespace is None else namespace
+    namespace[RUNTIME_NAME] = RUNTIME
+    exec(compile_source(source, filename), namespace)  # noqa: S102 - executing the caller's code is the purpose
+    return namespace