trace.py

"""
Trace Python execution at a function level.

Usage::

    import trace
    trace.start()

Currently only traces the current thread, but it could be extended to handle
multiple threads using :py:func:`threading.setprofile` and
:py:func:`threading.get_native_id()`.
"""
import sys
import inspect


def start():
    Tracer().register()


class Tracer:
    def __init__(self):
        self.depth = 0
        self.indent = "  "
        self.file = sys.stderr
        self.original_profiler = None

    def register(self, preserve = True):
        # Reset depth now since we'll start tracing right away with the return
        # of this method.
        self.depth = len(inspect.stack(context = 0)) - 1

        if preserve:
            self.original_profiler = sys.getprofile()

        sys.setprofile(self)
        # We'll start getting called here.

    def __call__(self, frame, event, arg):
        """
        Record a profile event.

        *frame* is a `frame object`_.

        *event* and *arg* are described by `sys.setprofile()`_.

        .. _frame object: https://docs.python.org/3/reference/datamodel.html#frame-objects
        .. _sys.setprofile(): https://docs.python.org/3/library/sys.html#sys.setprofile
        """
        if event == "call":
            # A function is called (or some other code block entered). The
            # profile function is called; arg is None.
            self.print(location(frame.f_back), "→ entering", function_name(frame))
            self.depth += 1

        elif event == "return":
            # A function (or other code block) is about to return. The profile
            # function is called; arg is the value that will be returned, or
            # None if the event is caused by an exception being raised.
            self.print(location(frame), "← exiting", function_name(frame), "=", repr(arg))
            self.depth -= 1

        elif event == "c_call":
            # A C function is about to be called. This may be an extension
            # function or a built-in. arg is the C function object.
            self.print(location(frame), "→ entering", f"C function {arg!r}")
            self.depth += 1

        elif event == "c_return":
            # A C function has returned. arg is the C function object.
            self.depth -= 1
            self.print(location(frame), "← exited", f"C function {arg!r}")

        elif event == "c_exception":
            # A C function has raised an exception. arg is the C function
            # object.
            self.depth -= 1
            self.print(location(frame), "← exited", f"C function {arg!r} due to exception")

        else:
            self.print(f"unknown trace event: {frame=} {event=} {arg=}")

        if self.original_profiler:
            self.original_profiler(frame, event, arg)

    def entering(self, *args):
        self.depth += 1

    def leaving(self, *args):
        self.depth -= 1

    def print(self, *args):
        print(self.indent * self.depth + " ".join(map(str, args)), file = self.file)


def function_name(frame):
    code = frame.f_code
    module = module_name(frame) or "<main>"
    return f"{module}.{code.co_name} ({code.co_filename}:{code.co_firstlineno})"


def module_name(frame):
    module = inspect.getmodule(frame.f_code)
    return module.__name__ if module else None


def location(frame):
    return f"{frame.f_code.co_filename}:{frame.f_lineno}" if frame else "hic sunt dracones"