A script to trace function calls in Python (for Debugging Purposes) and it saves locally to trace.log

debug_tracer.py

# USE IT LIKE THIS in main.py, for example:
"""
_TRACE_CALLS = install_call_tracing(app_root=ROOT, source_root=ROOT.parent)
"""
from __future__ import annotations

import atexit
import os
import sys
import time
from pathlib import Path
from types import FrameType


TRACE_IGNORED_DIRS = {
    ".git",
    ".mypy_cache",
    ".pytest_cache",
    ".ruff_cache",
    ".venv",
    ".env",
    "env",
    "venv",
    "__pycache__",
    "build",
    "dist",
    "htmlcov",
    "node_modules",
    "playwright-report",
    "test-results",
}


def install_call_tracing(*, app_root: Path, source_root: Path) -> CallTracer | None:
    if not _env_flag("DEBUG"):
        return None

    tracer = CallTracer(
        target_function=os.getenv("TRACE_FUNCTION", "event_stream").strip() or "event_stream",
        mode=os.getenv("TRACE_MODE", "project").strip().lower() or "project",
        output_path=_resolve_trace_output_path(app_root),
        colorize=_env_flag("TRACE_COLOR", default=True) and sys.stderr.isatty(),
        source_root=source_root.resolve(),
    )
    atexit.register(tracer.close)
    sys.settrace(tracer)
    return tracer


def _env_flag(name: str, default: bool = False) -> bool:
    raw_value = os.getenv(name)
    if raw_value is None:
        return default
    return raw_value.strip().lower() in {"1", "true", "yes", "on"}


def _resolve_trace_output_path(app_root: Path) -> Path:
    raw_value = os.getenv("TRACE_OUTPUT", "trace.log").strip() or "trace.log"
    output_path = Path(raw_value)
    if output_path.is_absolute():
        return output_path
    return app_root / output_path


def _is_project_source_file(file_name: str, *, source_root: Path) -> bool:
    if not file_name or file_name.startswith("<"):
        return False

    try:
        file_path = Path(file_name).resolve()
        file_path.relative_to(source_root)
    except (OSError, RuntimeError, ValueError):
        return False

    return not any(part in TRACE_IGNORED_DIRS for part in file_path.parts)


class CallTracer:
    _RESET = "\033[0m"
    _DIM = "\033[2m"
    _CYAN = "\033[36m"
    _MAGENTA = "\033[35m"
    _TREE = "\033[90m"

    def __init__(
        self,
        *,
        target_function: str,
        mode: str,
        output_path: Path,
        colorize: bool,
        source_root: Path,
    ) -> None:
        self.target_function = target_function
        self.mode = mode if mode in {"all", "project"} else "all"
        self.output_path = output_path
        self.colorize = colorize
        self.source_root = source_root
        self.tracking = False
        self.root_frame_id: int | None = None
        self.frame_stack: list[tuple[int, bool]] = []
        self.visible_frame_ids: list[int] = []
        self.log_file = self.output_path.open("w", encoding="utf-8")
        self._write_header()

    def __call__(self, frame: FrameType, event: str, arg: object):
        del arg
        if event == "call":
            self._handle_call(frame)
        elif event == "return":
            self._handle_return(frame)
        return self

    def close(self) -> None:
        if self.log_file.closed:
            return
        self.log_file.flush()
        self.log_file.close()

    def _handle_call(self, frame: FrameType) -> None:
        frame_id = id(frame)
        func_name = frame.f_code.co_name

        if not self.tracking:
            if func_name != self.target_function:
                return
            self.tracking = True
            self.root_frame_id = frame_id

        is_visible = self._should_log_frame(frame)
        self.frame_stack.append((frame_id, is_visible))
        if not is_visible:
            return

        visible_depth = len(self.visible_frame_ids)
        self._write_line(
            plain_text=self._plain_call_line(frame, visible_depth),
            colored_text=self._colored_call_line(frame, visible_depth),
        )
        self.visible_frame_ids.append(frame_id)

    def _handle_return(self, frame: FrameType) -> None:
        if not self.tracking:
            return

        frame_id = id(frame)
        is_visible = False
        if self.frame_stack and self.frame_stack[-1][0] == frame_id:
            _, is_visible = self.frame_stack.pop()

        if is_visible and self.visible_frame_ids and self.visible_frame_ids[-1] == frame_id:
            self.visible_frame_ids.pop()

        if frame_id == self.root_frame_id:
            self.tracking = False
            self.root_frame_id = None
            self.frame_stack.clear()
            self.visible_frame_ids.clear()

    def _should_log_frame(self, frame: FrameType) -> bool:
        if self.mode == "all":
            return True
        return _is_project_source_file(frame.f_code.co_filename, source_root=self.source_root)

    def _write_header(self) -> None:
        self.log_file.write(
            "TRACE CONFIG :: "
            f"target={self.target_function} "
            f"mode={self.mode} "
            f"output={self.output_path.name}\n"
        )
        self.log_file.write(f"TRACE ROOT :: {self.source_root}\n\n")
        self.log_file.flush()

    def _write_line(self, *, plain_text: str, colored_text: str) -> None:
        self.log_file.write(f"{plain_text}\n")
        self.log_file.flush()
        if self.colorize:
            print(colored_text, file=sys.stderr, flush=True)

    def _plain_call_line(self, frame: FrameType, depth: int) -> str:
        timestamp = time.strftime("%H:%M:%S")
        func_name = frame.f_code.co_name
        location = self._format_location(frame.f_code.co_filename, frame.f_lineno)
        tree_prefix = self._tree_prefix(depth)
        return f"{timestamp} | {tree_prefix}{func_name} ({location})"

    def _colored_call_line(self, frame: FrameType, depth: int) -> str:
        timestamp = time.strftime("%H:%M:%S")
        func_name = frame.f_code.co_name
        location = self._format_location(frame.f_code.co_filename, frame.f_lineno)
        tree_prefix = self._tree_prefix(depth)
        return (
            f"{self._DIM}{timestamp}{self._RESET} | "
            f"{self._TREE}{tree_prefix}{self._RESET}"
            f"{self._CYAN}{func_name}{self._RESET} "
            f"{self._MAGENTA}({location}){self._RESET}"
        )

    def _format_location(self, file_name: str, line_number: int) -> str:
        if _is_project_source_file(file_name, source_root=self.source_root):
            try:
                relative_path = Path(file_name).resolve().relative_to(self.source_root)
                return f"{relative_path}:{line_number}"
            except (OSError, RuntimeError, ValueError):
                pass
        return f"{file_name}:{line_number}"

    def _tree_prefix(self, depth: int) -> str:
        if depth == 0:
            return ""
        return f"{'│   ' * (depth - 1)}├── "