code cleanup

offline_debug/__init__.py

@@ -1,5 +1,6 @@
 """Tool for serializing and reconstructing Python exceptions with full stack traces."""
 
-from .serializer import load_traceback, save_traceback
+from ._inner.load_traceback import load_traceback
+from ._inner.save_traceback import save_traceback
 
 __all__ = ["load_traceback", "save_traceback"]

offline_debug/_inner/c_api/__init__.py

@@ -0,0 +1,4 @@
+from ._create_frame import create_frame
+from ._link_frame import link_frame
+
+__all__ = ["create_frame", "link_frame"]

offline_debug/_inner/c_api/_create_frame.py

@@ -0,0 +1,93 @@
+"""
+Abstraction of the CPython api for creating a traceback frame object.
+
+The reason we have to implement it, is that cpython doesn't expose
+the api to create a FrameType using pure python.
+"""
+
+from __future__ import annotations
+
+import ctypes
+import sys
+from functools import cache
+from types import CodeType, FrameType
+from typing import TYPE_CHECKING, Any
+
+if TYPE_CHECKING:
+    import _ctypes
+
+
+@cache
+def _get_py_frame_new() -> ctypes._NamedFuncPointer:
+    """Get the PyFrame_New C function configured with ctypes."""
+    func: ctypes._NamedFuncPointer = ctypes.pythonapi.PyFrame_New
+    func.argtypes = (
+        ctypes.c_void_p,  # PyThreadState *tstate
+        ctypes.py_object,  # PyCodeObject *code
+        ctypes.py_object,  # PyObject *globals
+        ctypes.py_object,  # PyObject *locals
+    )
+    func.restype = ctypes.py_object
+
+    def errcheck[T](
+        result: T | None,
+        _func: _ctypes.CFuncPtr,
+        _args: tuple[Any, ...],
+    ) -> T:  # pragma: no cover
+        if not result:
+            msg = "failed to create a new frame while calling"
+            raise RuntimeError(msg)
+        return result
+
+    func.errcheck = errcheck
+    return func
+
+
+@cache
+def _get_py_thread_state_get() -> ctypes._NamedFuncPointer:
+    """Get the PyThreadState_Get C function configured with ctypes."""
+    func = ctypes.pythonapi.PyThreadState_Get
+    func.argtypes = ()
+    func.restype = ctypes.c_void_p
+
+    def errcheck[T](
+        result: T | None, _func: _ctypes.CFuncPtr, _args: tuple[Any, ...]
+    ) -> T:  # pragma: no cover
+        if not result:
+            msg = "failed to get the current thread state"
+            raise RuntimeError(msg)
+        return result
+
+    func.errcheck = errcheck
+    return func
+
+
+def create_frame(
+    code: CodeType,
+    frame_globals: dict[str, Any],
+    frame_locals: dict[str, Any],
+    thread_state: int | None = None,
+) -> FrameType:
+    """
+    Create a FrameType object using the received parameters.
+
+    If the thread state received is None, use the current thread state automatically.
+    """
+    py_frame_new = _get_py_frame_new()
+    py_thread_state_get = _get_py_thread_state_get()
+
+    if thread_state is None:
+        thread_state = py_thread_state_get()
+
+    frame: FrameType = py_frame_new(thread_state, code, frame_globals, frame_locals)
+
+    # In 3.13+, PEP 667 allows safe write-through access to locals.
+    # py_frame_new ignores the frame_locals argument,
+    # and must be assigned after the frame's creation in 3.13+
+    if sys.version_info >= (3, 13) and frame_locals:
+        frame.f_locals.update(frame_locals)
+
+    if not isinstance(frame, FrameType):
+        msg = f"Expected types.FrameType, but got {type(frame).__name__}"
+        raise TypeError(msg)
+    return frame

offline_debug/_inner/c_api/_link_frame.py

@@ -0,0 +1,128 @@
+"""
+Link 2 frames together, so that a linked chain of frames will eventually create a traceback.
+
+The native cpython api does not expose a way to link frames together
+"""
+
+from __future__ import annotations
+
+import ctypes
+import sys
+from functools import cache
+from typing import TYPE_CHECKING, Any
+
+from ._create_frame import create_frame
+
+if TYPE_CHECKING:
+    import _ctypes
+    from types import FrameType
+
+
+@cache
+def _get_py_incref() -> ctypes._NamedFuncPointer:
+    """
+    Get the Py_IncRef C function configured with ctypes.
+
+    Py_IncRef increases the reference count of an object.
+    We call this function when we want to "own" an object,
+    meaning that the object won't be deleted while we are still using it.
+
+    In our case, we use the function to hold the f_back reference to link between frames.
+    """
+    func: ctypes._NamedFuncPointer = ctypes.pythonapi.Py_IncRef
+    func.argtypes = (ctypes.py_object,)
+    func.restype = None
+
+    def errcheck[T](
+        result: T,
+        _func: _ctypes.CFuncPtr,
+        _args: tuple[Any, ...],
+    ) -> T:  # pragma: no cover
+        if result is not None:
+            msg = f"Unexpected {result=}, expected None."
+            raise TypeError(msg)
+        return result
+
+    func.errcheck = errcheck
+    return func
+
+
+def link_frame(frame: FrameType, f_back: FrameType) -> None:
+    """Link a frame to its parent frame using the discovered offset."""
+    _f_back_offset = _get_f_back_offset()
+    if _f_back_offset is None:
+        msg = "Failed discovering the offset for the f_back property."
+        raise RuntimeError(msg)
+
+    # In Python, setting f_back means the child frame now owns a reference
+    # to the parent frame. We must increment the reference count of the
+    # parent to reflect this.
+    py_incref = _get_py_incref()
+    py_incref(f_back)
+
+    # Use ctypes to write the address of the back frame into the discovered offset.
+    # We find the address of the f_back property, by finding the start of the frame struct,
+    # Then moving to the f_back offset in the struct.
+    ptr = ctypes.c_void_p.from_address(id(frame) + _f_back_offset)
+    ptr.value = id(f_back)
+
+
+@cache
+def _get_f_back_offset() -> int | None:
+    """
+    Dynamically discover the memory offset of f_back in PyFrameObject.
+
+    The offset for the f_back can change between python versions and operating systems,
+    So we find the location of the f_back property dynamically.
+    The general idea of the algorithm is to create a mock frame,
+    And scanning its memory until we hit the f_back property field.
+    We know we hit the f_back property by
+    setting the value of each slot in the memeory to another frame, and checking if the
+    f_back property was populated after changing that memory location.
+    """
+    try:
+        # Compile an empty code object that we can use to create a frame.
+        code = compile("pass", "<discovery>", "exec")
+
+        # Create 2 frames that we can attempt linking between to find the f_back property.
+        frame = create_frame(code=code, frame_globals={}, frame_locals={})
+        f_back_frame = create_frame(code=code, frame_globals={}, frame_locals={})
+
+        ptr_size = ctypes.sizeof(ctypes.c_void_p)
+        jump_size = ptr_size
+
+        ref_count_size = ptr_size  # The amount of references to the object
+        type_object_size = ptr_size
+        frame_header_size = ref_count_size + type_object_size
+        start = frame_header_size  # The header does not contain the f_back property
+
+        frame_size = sys.getsizeof(frame)
+
+        # We scan the memory ptr_size bytes forward, so we decrease it from the end.
+        # Since we want the last memory slot to be scanned as well,
+        # we add +1 because the for loop is exclusive to the end.
+        end = frame_size - ptr_size + 1
+
+        # We start scanning after the PyObject header (refcnt + type).
+        for offset in range(start, end, jump_size):
+            candidate_f_back_address = id(frame) + offset
+            candidate_f_back_ptr = ctypes.c_ssize_t.from_address(candidate_f_back_address)
+
+            # f_back is initially NULL (0) in a newly created frame,
+            # since no other frames are linked to it.
+            if candidate_f_back_ptr.value != 0:
+                continue
+
+            candidate_f_back_ptr.value = id(f_back_frame)
+            # If reading f_back via Python now returns our target, we found it.
+            if frame.f_back is f_back_frame:
+                # Success, but we must restore 0 so we don't mess up refcounts
+                # when 'frame' is eventually garbage collected.
+                candidate_f_back_ptr.value = 0
+                return offset
+            # Restore to 0 if this wasn't the correct offset.
+            candidate_f_back_ptr.value = 0
+
+    except Exception:  # noqa: BLE001
+        return None
+    return None

offline_debug/_inner/load_traceback.py

@@ -0,0 +1,129 @@
+"""Load traceback object from a dump file."""
+
+import marshal
+import pickle
+import sys
+import types
+from io import BytesIO
+from pathlib import Path
+from types import CodeType
+from typing import Never
+
+from offline_debug._inner.c_api import (
+    create_frame,
+    link_frame,
+)
+from offline_debug._inner.models import (
+    _ExceptionData,
+    _FrameData,
+)
+
+
+def _reconstruct_exc_data(data: _ExceptionData) -> BaseException:
+    """
+    Recursively reconstruct an exception from its serialized data.
+
+    Note on Python Locals:
+    Python uses two ways to store local variables:
+    1. "Slow" locals: A dictionary used for module-level code and class definitions.
+    2. "Fast" locals: A fixed-size array used for functions. This is faster than
+       dictionary lookups because variables are accessed by index.
+
+    During reconstruction, we must explicitly synchronize these because PyFrame_New
+    does not automatically populate the "fast" locals array from a dictionary.
+    """
+    exc: BaseException = pickle.loads(data.exc_pickle)  # noqa: S301
+    if not isinstance(exc, BaseException):
+        msg = f"Expected BaseException, but got {type(exc).__name__}"
+        raise TypeError(msg)
+
+    reconstructed_frames: list[tuple[types.FrameType, _FrameData]] = []
+    for f_data in data.tb_frames:
+        code: CodeType = marshal.loads(f_data.code)  # noqa: S302
+
+        # In Python 3.11 and 3.12, accessing f_locals on a frame created via
+        # PyFrame_New for optimized code (functions) causes a segmentation fault
+        # because the internal 'fast' locals array is not initialized.
+        # As a workaround, we create a 'non-optimized' version of the code object
+        # by compiling a dummy string. This ensures the bytecode is safe
+        # (no LOAD_FAST) while preserving metadata like name and filename.
+        if sys.version_info < (3, 13):
+            # A simple module-level code object never has fast locals.
+            # Since the source is empty, no optimized locals will be created.
+            # Instead, python will go to the unoptimized dictionary we set under frame_locals later.
+            unoptimized_code = compile("", code.co_filename, "exec")
+            code = unoptimized_code.replace(
+                co_name=code.co_name,
+                co_firstlineno=code.co_firstlineno,
+                co_qualname=code.co_qualname,
+            )
+
+        # PyFrame_New returns a new reference to a PyFrameObject.
+        if f_data.module_name:
+            f_data.globals["__name__"] = f_data.module_name
+
+        frame: types.FrameType = create_frame(
+            code=code, frame_globals=f_data.globals, frame_locals=f_data.locals
+        )
+
+        if reconstructed_frames:
+            # link the frame back to the previously constructed frame.
+            link_frame(frame, reconstructed_frames[-1][0])
+
+        reconstructed_frames.append((frame, f_data))
+
+    tb_next: types.TracebackType | None = None
+    for frame, f_data in reversed(reconstructed_frames):
+        tb = types.TracebackType(
+            tb_next=tb_next,
+            tb_frame=frame,
+            tb_lasti=f_data.lasti,
+            tb_lineno=f_data.lineno,
+        )
+        tb_next = tb
+
+    exc = exc.with_traceback(tb_next)
+
+    if data.cause:
+        exc.__cause__ = _reconstruct_exc_data(data.cause)
+    if data.context:
+        exc.__context__ = _reconstruct_exc_data(data.context)
+
+    return exc
+
+
+def load_traceback(file: Path | BytesIO) -> Never:
+    """Load an exception and its traceback from a file and raise it."""
+    if isinstance(file, Path):
+        with file.open("rb") as f:
+            data = pickle.load(f)  # noqa: S301
+    else:
+        data = pickle.load(file)  # noqa: S301
+
+    if not isinstance(data, _ExceptionData):
+        msg = f"Expected _ExceptionData, but got {type(data).__name__}"
+        raise TypeError(msg)
+
+    exc = _reconstruct_exc_data(data)
+
+    current_frames: list[types.FrameType] = []
+    curr: types.FrameType | None = sys._getframe(1)  # noqa: SLF001
+    while curr:
+        current_frames.append(curr)
+        curr = curr.f_back
+
+    if exc.__traceback__ and current_frames:
+        reconstructed_outer = exc.__traceback__.tb_frame
+        link_frame(reconstructed_outer, current_frames[0])
+
+    tb_chain: types.TracebackType | None = exc.__traceback__
+    for frame in current_frames:
+        tb_chain = types.TracebackType(
+            tb_next=tb_chain,
+            tb_frame=frame,
+            tb_lasti=frame.f_lasti,
+            tb_lineno=frame.f_lineno,
+        )
+
+    exc = exc.with_traceback(tb_chain)
+    raise exc

offline_debug/_inner/models.py

@@ -0,0 +1,29 @@
+"""Functions for serializing and reconstructing exceptions with their tracebacks."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Any
+
+
+@dataclass
+class _FrameData:
+    """Serialized data for a single stack frame."""
+
+    code: bytes
+    globals: dict[str, Any]
+    locals: dict[str, Any]
+    lasti: int
+    lineno: int
+    stack_depth: int
+    module_name: str | None = None
+
+
+@dataclass
+class _ExceptionData:
+    """Serialized data for an exception and its traceback."""
+
+    exc_pickle: bytes
+    tb_frames: list[_FrameData]
+    cause: _ExceptionData | None = None
+    context: _ExceptionData | None = None