tracer.py

"""Custom sys.settrace-based execution tracer for the Python Code Visualizer."""
import io
import sys
import types
from collections import deque

MAX_ITEMS = 50
MAX_DEPTH = 20
MAX_STEPS = 1000
USER_CODE_FILENAME = "<user_code>"


class _StepLimitReached(BaseException):
    # Subclasses BaseException (not Exception) on purpose: user code containing
    # `except Exception:` or a bare `except:` must not be able to swallow the
    # sentinel that halts execution at the step cap. Same technique CPython
    # uses for KeyboardInterrupt/SystemExit.
    pass


def serialize_value(value, depth=0, seen=None):
    if seen is None:
        seen = frozenset()

    if value is None or isinstance(value, (bool, int, float, str)):
        return value

    obj_id = id(value)

    if isinstance(value, (list, tuple)):
        if obj_id in seen:
            return {"type": "circular_ref"}
        if depth >= MAX_DEPTH:
            return {"type": "truncated_depth"}
        child_seen = seen | {obj_id}
        total = len(value)
        items = [serialize_value(v, depth + 1, child_seen) for v in list(value)[:MAX_ITEMS]]
        result = {"type": "list", "items": items}
        if total > MAX_ITEMS:
            result["truncated"] = True
            result["total_count"] = total
        return result

    if isinstance(value, dict):
        if obj_id in seen:
            return {"type": "circular_ref"}
        if depth >= MAX_DEPTH:
            return {"type": "truncated_depth"}
        child_seen = seen | {obj_id}
        total = len(value)
        pairs = []
        for k, v in list(value.items())[:MAX_ITEMS]:
            pairs.append([serialize_value(k, depth + 1, child_seen), serialize_value(v, depth + 1, child_seen)])
        result = {"type": "dict", "pairs": pairs}
        if total > MAX_ITEMS:
            result["truncated"] = True
            result["total_count"] = total
        return result

    if hasattr(value, "__dict__"):
        if obj_id in seen:
            return {"type": "circular_ref"}
        if depth >= MAX_DEPTH:
            return {"type": "truncated_depth"}
        child_seen = seen | {obj_id}

        if _looks_like_linked_list(value):
            return _serialize_linked_list(value, depth, child_seen)
        if _looks_like_tree(value):
            return _serialize_tree(value, depth, seen)

        attrs = {}
        for k, v in vars(value).items():
            if k.startswith("__"):
                continue
            attrs[k] = serialize_value(v, depth + 1, child_seen)
        return {"type": "object", "class": type(value).__name__, "attrs": attrs}

    return repr(value)


def _looks_like_linked_list(value):
    return hasattr(value, "next") and not (hasattr(value, "left") or hasattr(value, "right"))


def _node_value(node, exclude, depth, seen):
    attrs = {k: v for k, v in vars(node).items() if k not in exclude and not k.startswith("__")}
    if len(attrs) == 1:
        return serialize_value(next(iter(attrs.values())), depth + 1, seen)
    return {k: serialize_value(v, depth + 1, seen) for k, v in attrs.items()}


def _serialize_linked_list(value, depth, seen):
    nodes = []
    visited = set()
    total = 0
    current = value
    while current is not None:
        if id(current) in visited:
            break
        visited.add(id(current))
        total += 1
        if len(nodes) < MAX_ITEMS:
            nodes.append(_node_value(current, {"next"}, depth, seen))
        current = getattr(current, "next", None)
    result = {"type": "linked_list", "nodes": nodes}
    if total > MAX_ITEMS:
        result["truncated"] = True
        result["total_count"] = total
    return result


def _looks_like_tree(value):
    return hasattr(value, "left") and hasattr(value, "right")


def _serialize_tree(value, depth, seen):
    included = set()
    visited = set()
    total = 0
    queue = deque([value])
    while queue:
        node = queue.popleft()
        if node is None or id(node) in visited:
            continue
        visited.add(id(node))
        total += 1
        if len(included) < MAX_ITEMS:
            included.add(id(node))
        queue.append(getattr(node, "left", None))
        queue.append(getattr(node, "right", None))

    def build(node, d):
        if node is None:
            return None
        if id(node) in seen:
            return {"type": "circular_ref"}
        if d >= MAX_DEPTH:
            return {"type": "truncated_depth"}
        if id(node) not in included:
            return {"type": "truncated_marker"}
        return {
            "val": _node_value(node, {"left", "right"}, d, seen),
            "left": build(getattr(node, "left", None), d + 1),
            "right": build(getattr(node, "right", None), d + 1),
        }

    root = build(value, depth)
    result = {"type": "tree", "root": root}
    if total > MAX_ITEMS:
        result["truncated"] = True
        result["total_count"] = total
    return result


def _is_skippable_local(value):
    return isinstance(value, (types.FunctionType, types.ModuleType, type, types.BuiltinFunctionType))


def run_and_trace(code_str):
    try:
        compiled = compile(code_str, USER_CODE_FILENAME, "exec")
    except SyntaxError as e:
        return {
            "steps": [],
            "stdout": "",
            "error": f"SyntaxError: {e.msg} (line {e.lineno})",
            "step_limit_reached": False,
        }

    steps = []
    call_ids = {}
    call_id_counter = [0]
    stop = [False]

    def capture_frames(top_frame):
        chain = []
        f = top_frame
        while f is not None and id(f) in call_ids:
            chain.append(f)
            f = f.f_back
        chain.reverse()
        result = []
        for fr in chain:
            locals_out = {}
            for k, v in fr.f_locals.items():
                if k.startswith("__") or _is_skippable_local(v):
                    continue
                locals_out[k] = serialize_value(v)
            name = fr.f_code.co_name
            result.append({
                "call_id": call_ids[id(fr)],
                "name": "module" if name == "<module>" else name,
                "locals": locals_out,
                "line": fr.f_lineno,
            })
        return result

    def record_step(frame, event):
        if stop[0]:
            return
        step = {
            "step": len(steps),
            "line": frame.f_lineno,
            "event": event,
            "frames": capture_frames(frame),
        }
        steps.append(step)
        if len(steps) >= MAX_STEPS:
            stop[0] = True
            sys.settrace(None)
            raise _StepLimitReached()

    exception_frames = set()

    def local_trace(frame, event, arg):
        if stop[0]:
            return None
        if event == "line":
            exception_frames.discard(id(frame))
            record_step(frame, "line")
        elif event == "return":
            if id(frame) in exception_frames:
                exception_frames.discard(id(frame))
            else:
                record_step(frame, "return")
                call_ids.pop(id(frame), None)
        elif event == "exception":
            exception_frames.add(id(frame))
        return None if stop[0] else local_trace

    def global_trace(frame, event, arg):
        if event != "call" or stop[0]:
            return None
        if frame.f_code.co_filename != USER_CODE_FILENAME:
            return None
        call_id_counter[0] += 1
        call_ids[id(frame)] = call_id_counter[0]
        record_step(frame, "call")
        return None if stop[0] else local_trace

    stdout_capture = io.StringIO()
    old_stdout = sys.stdout
    namespace = {"__name__": "__main__"}
    sys.stdout = stdout_capture
    sys.settrace(global_trace)
    try:
        exec(compiled, namespace)
    except _StepLimitReached:
        pass
    except BaseException as e:
        tb = e.__traceback__
        frame_chain = []
        while tb is not None:
            f = tb.tb_frame
            if f.f_code.co_filename == USER_CODE_FILENAME:
                frame_chain.append(f)
            tb = tb.tb_next
        frames_info = []
        for f in frame_chain:
            locals_out = {}
            for k, v in f.f_locals.items():
                if k.startswith("__") or _is_skippable_local(v):
                    continue
                locals_out[k] = serialize_value(v)
            name = f.f_code.co_name
            frames_info.append({
                "call_id": call_ids.get(id(f), 0),
                "name": "module" if name == "<module>" else name,
                "locals": locals_out,
                "line": f.f_lineno,
            })
        steps.append({
            "step": len(steps),
            "line": frame_chain[-1].f_lineno if frame_chain else 0,
            "event": "exception",
            "frames": frames_info,
            "exception": {"type": type(e).__name__, "message": str(e)},
        })
    finally:
        sys.settrace(None)
        sys.stdout = old_stdout

    last = steps[-1] if steps else None
    error = None
    if last and last["event"] == "exception":
        error = f"{last['exception']['type']}: {last['exception']['message']}"

    return {
        "steps": steps,
        "stdout": stdout_capture.getvalue(),
        "error": error,
        "step_limit_reached": len(steps) >= MAX_STEPS,
    }