jupyter-python-tutor 0.2.0__py3-none-any.whl

jupyter_python_tutor/tracer.py

@@ -0,0 +1,591 @@
+"""Code tracer using sys.settrace() to capture step-by-step execution state.
+
+Produces a list of TraceStep dicts, each containing the line number, the
+source line text, and the state of all local/global variables at that point.
+
+Key design: the tracer executes code directly in the kernel's globals so
+that imported packages and variables from previous cells are available.
+Variables defined by the traced code persist in the kernel (like a normal
+Jupyter cell).  Each object that is too complex to show inline gets a
+stable ``heap_id`` label (``id1``, ``id2``, ...) so the renderer can draw
+arrow boxes from frame variables to heap objects.
+"""
+from __future__ import annotations
+
+import ast
+import re as _re
+import sys
+import types
+import typing
+from copy import deepcopy
+
+__all__ = ["trace_code"]
+
+
+# Types we consider "simple" — shown inline in frames.
+_SIMPLE = (int, float, str, bool, type(None))
+_CONTAINER = (list, tuple, set, frozenset, dict)
+
+
+def _is_simple(obj: typing.Any) -> bool:
+    """Return True if obj is a primitive we can show inline."""
+    if isinstance(obj, _SIMPLE):
+        return True
+    # Small containers of primitives can also be inline
+    if isinstance(obj, _CONTAINER):
+        try:
+            items = obj if not isinstance(obj, dict) else list(obj.values())
+            if len(items) > 8:
+                return False
+            return all(_is_simple(i) for i in items)
+        except Exception:
+            return False
+    return False
+
+
+def _is_visualizable(obj: typing.Any) -> bool:
+    """Return True if obj is a type we can safely snapshot for visualization."""
+    if isinstance(obj, _SIMPLE):
+        return True
+    if isinstance(obj, _CONTAINER):
+        try:
+            for item in obj if not isinstance(obj, dict) else obj.values():
+                if not _is_visualizable(item):
+                    return False
+            return True
+        except Exception:
+            return False
+    return False
+
+
+def _safe_repr(obj: typing.Any) -> str:
+    """Return a repr, falling back to type name for non-visualizable objects."""
+    try:
+        r = repr(obj)
+        if len(r) > 200:
+            r = r[:197] + "..."
+        return r
+    except Exception:
+        return f"<{type(obj).__name__}>"
+
+
+# ─── Names to exclude from the global frame visualization ──────────────────
+# These are IPython/Jupyter internals and Python builtins that clutter the
+# Global frame.  We want the Global frame to show only user-relevant names:
+# imported modules (as references, not contents), user-defined variables,
+# classes, and functions defined in the notebook or the traced cell.
+_EXCLUDE_GLOBAL_NAMES = frozenset({
+    # IPython internals
+    "__name__", "__doc__", "__package__", "__loader__", "__spec__",
+    "__builtins__", "__file__", "__cached__", "_dh", "_ih", "_oh",
+    "_sh", "_", "__", "___", "__builtin__", "__builtins__",
+    "In", "Out", "get_ipython", "exit", "quit", "user_ns",
+    "open", "input",  # builtins that shouldn't appear as global vars
+    "_i", "_ii", "_iii", "_i1", "_i2", "_i3", "_i4", "_i5",
+    "_i6", "_i7", "_i8", "_i9", "_i10", "_i11", "_i12",
+    # Jupyter/IPython helpers
+    "display", "HTML", "JSON", "Markdown", "Math", "Latex",
+    "PrettyPrinter", "pretty",
+})
+
+# Names that look like IPython history cells (_1, _2, etc.)
+_IPYTHON_HISTORY = _re.compile(r"^_\d+$")
+
+# Module-level cache: maps code_str → set of names the code touches.
+# This avoids re-parsing on every trace step.
+_code_name_cache: dict[str, frozenset[str]] = {}
+
+
+def _extract_code_names(code_str: str) -> frozenset[str]:
+    """Extract all names the traced code defines or references via AST.
+
+    Returns a frozenset of variable names, function names, class names,
+    and imported names that appear in the code.  This is used to filter
+    the global frame so only code-relevant variables are shown.
+    """
+    if code_str in _code_name_cache:
+        return _code_name_cache[code_str]
+
+    names: set[str] = set()
+    try:
+        tree = ast.parse(code_str)
+        for node in ast.walk(tree):
+            if isinstance(node, ast.Name):
+                names.add(node.id)
+            elif isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef, ast.ClassDef)):
+                names.add(node.name)
+            elif isinstance(node, ast.arg):
+                names.add(node.arg)
+            elif isinstance(node, ast.alias):
+                # import statements: "import os" → "os", "from m import x as y" → "y"
+                names.add(node.asname or node.name)
+            elif isinstance(node, ast.Global):
+                for n in node.names:
+                    names.add(n)
+            elif isinstance(node, ast.Nonlocal):
+                for n in node.names:
+                    names.add(n)
+    except SyntaxError:
+        pass
+
+    result = frozenset(names)
+    _code_name_cache[code_str] = result
+    return result
+
+
+def _is_excluded_global(name: str, val: typing.Any) -> bool:
+    """Return True if a global variable should NOT be shown in the Global frame."""
+    if name in _EXCLUDE_GLOBAL_NAMES:
+        return True
+    if name.startswith("__"):
+        return True
+    if _IPYTHON_HISTORY.match(name):
+        return True
+    # Skip IPython display hooks
+    if name.startswith("_") and name not in ("_", "__", "___"):
+        # Allow _, __, ___ (last three results) but skip other _ prefixed
+        if not name in ("_", "__", "___"):
+            return True
+    return False
+
+
+def _capture_object_details(val: typing.Any, id_map: dict[int, str]) -> dict:
+    """Capture details about an object for the heap display.
+    
+    Returns a dict with:
+      - attrs: list of {name, value_repr, heap_id} for instance/class attributes
+      - items: list of {key, value_repr, heap_id} for containers
+      - bases: list of base class names (for classes)
+    """
+    details: dict = {"attrs": [], "items": [], "bases": []}
+    
+    try:
+        if isinstance(val, type):
+            # Class object — capture base classes and methods
+            bases = [b.__name__ for b in val.__bases__ if b is not object]
+            details["bases"] = bases
+            # Capture methods defined on this class (not inherited)
+            for name in sorted(dir(val)):
+                if name.startswith("__") and name != "__init__":
+                    continue
+                try:
+                    member = getattr(val, name)
+                    if callable(member) and not isinstance(member, type):
+                        details["attrs"].append({"name": name, "value_repr": "", "heap_id": None})
+                except Exception:
+                    pass
+        elif isinstance(val, (types.FunctionType, types.BuiltinFunctionType)):
+            # Function — capture __name__ and __qualname__
+            # IMPORTANT: functools.wraps copies __name__ from the wrapped function
+            # to the wrapper, so val.__name__ may be "fibonacci" even for a wrapper.
+            # The true function name is in val.__code__.co_name (e.g. "wrapper").
+            # We show __name__ as an attribute so the user can see the renaming.
+            details["attrs"].append({"name": "__name__", "value_repr": val.__name__, "heap_id": None})
+            if hasattr(val, "__qualname__"):
+                details["attrs"].append({"name": "__qualname__", "value_repr": val.__qualname__, "heap_id": None})
+        elif isinstance(val, types.ModuleType):
+            # Module — just the name
+            pass
+        elif isinstance(val, dict):
+            # Dict — capture key-value pairs
+            for k, v in list(val.items())[:12]:
+                if _is_simple(v):
+                    details["items"].append({"key": _safe_repr(k), "value_repr": _safe_repr(v), "heap_id": None})
+                else:
+                    oid = id(v)
+                    if oid not in id_map:
+                        id_map[oid] = f"id{len(id_map) + 1}"
+                    details["items"].append({"key": _safe_repr(k), "value_repr": "", "heap_id": id_map[oid]})
+        elif isinstance(val, (list, tuple, set, frozenset)):
+            # Sequence — capture items
+            for i, v in enumerate(list(val)[:12]):
+                if _is_simple(v):
+                    details["items"].append({"key": str(i), "value_repr": _safe_repr(v), "heap_id": None})
+                else:
+                    oid = id(v)
+                    if oid not in id_map:
+                        id_map[oid] = f"id{len(id_map) + 1}"
+                    details["items"].append({"key": str(i), "value_repr": "", "heap_id": id_map[oid]})
+        else:
+            # Instance object — capture instance attributes
+            if hasattr(val, "__dict__"):
+                for name, v in sorted(vars(val).items()):
+                    if name.startswith("__"):
+                        continue
+                    if _is_simple(v):
+                        details["attrs"].append({"name": name, "value_repr": _safe_repr(v), "heap_id": None})
+                    else:
+                        oid = id(v)
+                        if oid not in id_map:
+                            id_map[oid] = f"id{len(id_map) + 1}"
+                        details["attrs"].append({"name": name, "value_repr": "", "heap_id": id_map[oid]})
+    except Exception:
+        pass
+    
+    return details
+
+
+def _snapshot_variables(frame: types.FrameType, id_map: dict[int, str] | None = None) -> tuple[dict, list[dict]]:
+    """Capture local variables and return (simple_vars, object_refs).
+
+    simple_vars: name→repr for primitives/small containers (shown inline)
+    object_refs: list of {name, repr, id, heap_id, type} for objects that should be
+                 shown as heap references (classes, instances, large containers)
+    id_map: shared dict mapping id(obj)→heap_id label, so the same object
+            always gets the same label across all frames and steps.
+    """
+    if id_map is None:
+        id_map = {}
+    simple: dict[str, str] = {}
+    objects: list[dict] = []
+
+    # Locals (function frame)
+    for name, val in frame.f_locals.items():
+        if name.startswith("__"):
+            continue
+        if _is_simple(val):
+            simple[name] = _safe_repr(val)
+        else:
+            oid = id(val)
+            if oid not in id_map:
+                id_map[oid] = f"id{len(id_map) + 1}"
+            objects.append({
+                "name": name,
+                "repr": _safe_repr(val),
+                "id": oid,
+                "heap_id": id_map[oid],
+                "type": type(val).__name__,
+                "is_class": isinstance(val, type),
+                "is_instance": hasattr(val, "__class__") and not isinstance(val, type) and not isinstance(val, _SIMPLE + _CONTAINER + (types.FunctionType, types.ModuleType)),
+                "is_function": isinstance(val, (types.FunctionType, types.BuiltinFunctionType)),
+                "is_module": isinstance(val, types.ModuleType),
+                # true_name: the actual function name from the code object.
+                # functools.wraps copies __name__ but __code__.co_name stays original.
+                "true_name": val.__code__.co_name if isinstance(val, types.FunctionType) else None,
+                "display_name": val.__name__ if isinstance(val, types.FunctionType) else None,
+                "details": _capture_object_details(val, id_map),
+            })
+
+    return simple, objects
+
+
+def _snapshot_globals(
+    frame: types.FrameType,
+    id_map: dict[int, str] | None = None,
+    code_names: frozenset[str] | None = None,
+) -> tuple[dict, list[dict]]:
+    """Capture user globals (not builtins/IPython internals).
+
+    Returns (simple_vars, object_refs) same as _snapshot_variables but
+    for the global frame, filtered to exclude clutter.
+
+    If ``code_names`` is provided, only globals whose name appears in
+    ``code_names`` are shown — this limits the display to variables and
+    imports that the traced code actually uses.
+    """
+    if id_map is None:
+        id_map = {}
+    simple: dict[str, str] = {}
+    objects: list[dict] = []
+
+    for name, val in frame.f_globals.items():
+        if _is_excluded_global(name, val):
+            continue
+        # If we have code-relevant names, filter to only those
+        if code_names is not None and name not in code_names:
+            continue
+        # Skip modules — they're imports, show as references
+        if isinstance(val, types.ModuleType):
+            # Show imported modules as references (like OPT_Mentor shows
+            # imported packages as arrows to module objects)
+            oid = id(val)
+            if oid not in id_map:
+                id_map[oid] = f"id{len(id_map) + 1}"
+            objects.append({
+                "name": name,
+                "repr": f"<module '{val.__name__}'>",
+                "id": oid,
+                "heap_id": id_map[oid],
+                "type": "module",
+                "is_class": False,
+                "is_instance": False,
+                "is_function": False,
+                "is_module": True,
+                "details": _capture_object_details(val, id_map),
+            })
+            continue
+        if _is_simple(val):
+            simple[name] = _safe_repr(val)
+        else:
+            oid = id(val)
+            if oid not in id_map:
+                id_map[oid] = f"id{len(id_map) + 1}"
+            objects.append({
+                "name": name,
+                "repr": _safe_repr(val),
+                "id": oid,
+                "heap_id": id_map[oid],
+                "type": type(val).__name__,
+                "is_class": isinstance(val, type),
+                "is_instance": hasattr(val, "__class__") and not isinstance(val, type) and not isinstance(val, _SIMPLE + _CONTAINER + (types.FunctionType, types.ModuleType)),
+                "is_function": isinstance(val, (types.FunctionType, types.BuiltinFunctionType)),
+                "is_module": isinstance(val, types.ModuleType),
+                "true_name": val.__code__.co_name if isinstance(val, types.FunctionType) else None,
+                "display_name": val.__name__ if isinstance(val, types.FunctionType) else None,
+                "details": _capture_object_details(val, id_map),
+            })
+
+    return simple, objects
+
+
+def _snapshot_frame(
+    frame: types.FrameType,
+    is_global: bool = False,
+    id_map: dict[int, str] | None = None,
+    code_names: frozenset[str] | None = None,
+) -> dict:
+    """Capture a single frame's name, line, and variables for the call stack."""
+    if is_global:
+        simple_vars, object_refs = _snapshot_globals(frame, id_map, code_names)
+    else:
+        simple_vars, object_refs = _snapshot_variables(frame, id_map)
+    return {
+        "func": frame.f_code.co_name,
+        "line": frame.f_lineno,
+        "simple_vars": simple_vars,
+        "object_refs": object_refs,
+        "locals": simple_vars,  # backward compat
+    }
+
+
+def _snapshot_call_stack(
+    frame: types.FrameType,
+    id_map: dict[int, str] | None = None,
+    code_names: frozenset[str] | None = None,
+) -> list[dict]:
+    """Walk up the call stack via frame.f_back and capture all frames.
+
+    Returns a list ordered outermost-first (module → outer → inner).
+    The first element is always the Global frame (using f_globals).
+    Only frames from our traced file are included (plus the global frame).
+    """
+    stack: list[dict] = []
+    f = frame
+    seen_globals = False
+    while f is not None:
+        if f.f_code.co_filename == frame.f_code.co_filename:
+            is_global = (f.f_code.co_name == "<module>")
+            if is_global:
+                seen_globals = True
+            stack.append(_snapshot_frame(f, is_global=is_global, id_map=id_map, code_names=code_names))
+        f = f.f_back
+    stack.reverse()  # outermost first
+
+    # If we didn't see a global frame in the traced file, add one from f_globals
+    if not seen_globals and frame.f_globals:
+        # Create a synthetic global frame
+        simple_vars, object_refs = _snapshot_globals(frame, id_map, code_names)
+        global_frame = {
+            "func": "<module>",
+            "line": 0,
+            "simple_vars": simple_vars,
+            "object_refs": object_refs,
+            "locals": simple_vars,
+        }
+        stack.insert(0, global_frame)
+
+    return stack
+
+
+def _trace_code_internal(
+    code_str: str,
+    user_globals: dict[str, typing.Any] | None = None,
+    max_steps: int = 5000,
+) -> list[dict]:
+    """Execute *code_str* under a trace and return the list of steps.
+
+    Each step is a dict with keys:
+        line: int          — line number (1-based) that is about to execute
+        source: str        — the source line text (stripped)
+        call_stack: list   — list of frame dicts (outermost first)
+        event: str         — "call", "line", or "return"
+        func: str          — current function name
+        stdout: str        — output printed up to THIS step (incremental)
+    """
+    steps: list[dict] = []
+    step_count = 0
+    # Execute directly in the user's namespace so that classes, functions,
+    # and imports from previous cells are available. This is how a normal
+    # Jupyter cell works — the %%tutor magic should behave the same way,
+    # just with tracing enabled. Variables defined by the traced code WILL
+    # persist in the kernel (like a normal cell), but we must not overwrite
+    # __name__ or __builtins__.
+    if user_globals is None:
+        user_globals = {"__name__": "__main__", "__builtins__": __builtins__}
+    g = user_globals
+    # Ensure __builtins__ is available
+    if "__builtins__" not in g:
+        g["__builtins__"] = __builtins__
+
+    # Compile with a unique filename so we can map line numbers
+    filename = "<jpt>"
+    try:
+        compiled = compile(code_str, filename, "exec")
+    except SyntaxError as e:
+        return [{
+            "line": e.lineno or 1,
+            "source": f"# SyntaxError: {e.msg}",
+            "call_stack": [],
+            "event": "exception",
+            "func": "<module>",
+            "error": f"{type(e).__name__}: {e.msg}",
+            "stdout": "",
+        }]
+
+    source_lines = code_str.splitlines()
+
+    # Extract names the code touches, so we can filter the global frame
+    # to only show variables/imports relevant to this code.
+    code_names = _extract_code_names(code_str)
+
+    class _StepLimitExceeded(Exception):
+        """Raised internally to stop execution when max_steps is reached."""
+
+    # ─── Stdout capture with per-step tracking ──────────────────────────
+    # We capture stdout in a buffer and track the length at each step,
+    # so each step shows only the output produced up to that point.
+    import io
+    captured = io.StringIO()
+    stdout_lengths: list[int] = []  # length of captured output at each step
+
+    class _StepTrackingStream:
+        """A stream that tracks how much output has been written at each step."""
+        def __init__(self, real_stream):
+            self._real = real_stream
+            self._total_len = 0
+        def write(self, s):
+            self._real.write(s)
+            self._total_len += len(s)
+        def flush(self):
+            self._real.flush()
+        @property
+        def total_len(self):
+            return self._total_len
+
+    tracking_stream = _StepTrackingStream(captured)
+
+    # Shared ID map: maps id(obj) → "id1", "id2", etc.
+    # Same object always gets the same label across all steps.
+    id_map: dict[int, str] = {}
+
+    def _trace(frame, event, arg):
+        nonlocal step_count
+        if frame.f_code.co_filename != filename:
+            return _trace  # skip frames from other files
+        if step_count >= max_steps:
+            raise _StepLimitExceeded()  # stop execution, not just tracing
+        step_count += 1
+
+        lineno = frame.f_lineno
+        source = source_lines[lineno - 1].strip() if 1 <= lineno <= len(source_lines) else ""
+        call_stack = _snapshot_call_stack(frame, id_map, code_names)
+        current_stdout = captured.getvalue()  # everything printed so far
+
+        if event == "line":
+            steps.append({
+                "line": lineno,
+                "source": source,
+                "call_stack": call_stack,
+                "event": "line",
+                "func": frame.f_code.co_name,
+                "stdout": current_stdout,
+            })
+        elif event == "call":
+            steps.append({
+                "line": lineno,
+                "source": source,
+                "call_stack": call_stack,
+                "event": "call",
+                "func": frame.f_code.co_name,
+                "stdout": current_stdout,
+            })
+        elif event == "return":
+            ret_repr = _safe_repr(arg) if arg is not None else "None"
+            steps.append({
+                "line": lineno,
+                "source": source,
+                "call_stack": call_stack,
+                "event": "return",
+                "func": frame.f_code.co_name,
+                "return_value": ret_repr,
+                "stdout": current_stdout,
+            })
+        elif event == "exception":
+            exc_type, exc_val, _ = arg
+            # Don't record our internal exception
+            if exc_type is _StepLimitExceeded:
+                return None
+            steps.append({
+                "line": lineno,
+                "source": source,
+                "call_stack": call_stack,
+                "event": "exception",
+                "func": frame.f_code.co_name,
+                "error": f"{exc_type.__name__}: {exc_val}",
+                "stdout": current_stdout,
+            })
+
+        return _trace
+
+    old_stdout = sys.stdout
+    old_trace = sys.gettrace()
+
+    try:
+        sys.stdout = tracking_stream
+        sys.settrace(_trace)
+        exec(compiled, g)
+    except _StepLimitExceeded:
+        pass  # execution stopped due to step limit
+    except Exception as e:
+        steps.append({
+            "line": 0,
+            "source": f"# {type(e).__name__}: {e}",
+            "call_stack": [],
+            "event": "exception",
+            "func": "<module>",
+            "error": f"{type(e).__name__}: {e}",
+            "stdout": captured.getvalue(),
+        })
+    finally:
+        sys.settrace(old_trace)
+        sys.stdout = old_stdout
+
+    return steps
+
+
+def trace_code(
+    code_str: str,
+    max_steps: int = 5000,
+    user_globals: dict[str, typing.Any] | None = None,
+) -> list[dict]:
+    """Trace Python code and return step-by-step execution state.
+
+    Parameters
+    ----------
+    code_str : str
+        The Python code to trace.
+    max_steps : int
+        Maximum number of trace steps to capture (safety limit).
+    user_globals : dict, optional
+        The kernel's global namespace (so imported packages and existing
+        variables are available to the traced code).  The code is executed
+        directly in this namespace — variables defined by the traced code
+        persist in the kernel (like a normal Jupyter cell).
+
+    Returns
+    -------
+    list[dict]
+        Each dict has keys: line, source, call_stack, event, func.
+        May also include: return_value, error, stdout.
+    """
+    return _trace_code_internal(code_str, user_globals, max_steps)