Python - Simple call tracing profiler for Selenium projects

simple_selenium_profiler.py

import atexit
import importlib
import pkgutil
import sys
import time
from collections import defaultdict
from pathlib import Path

import selenium

PROJECT_ROOT = Path(__file__).parent.resolve()  # project code
TRACK_PACKAGES = ["selenium"]  # start with selenium


def collect_package_dirs(packages):
    """Find all directories containing specified packages or their submodules."""
    dirs = set()
    visited = set()
    stack = list(packages)
    while stack:
        pkg_name = stack.pop()
        if pkg_name in visited:
            continue
        visited.add(pkg_name)
        try:
            pkg = importlib.import_module(pkg_name)
        except Exception:
            continue
        # add package directory if available
        if hasattr(pkg, "__file__") and pkg.__file__:
            dirs.add(Path(pkg.__file__).parent.resolve())
        # add submodules for packages
        if hasattr(pkg, "__path__"):
            for finder, name, ispkg in pkgutil.iter_modules(pkg.__path__):
                full_name = pkg_name + "." + name
                stack.append(full_name)
    return dirs


PACKAGE_DIRS = collect_package_dirs(TRACK_PACKAGES)


stats = defaultdict(lambda: {"calls": 0, "total_time": 0.0})
frame_stack = {}


def trace(frame, event, arg):
    if event not in ("call", "return"):
        return trace
    filename = frame.f_code.co_filename
    if not filename or filename.startswith("<"):
        return trace  # skip built-ins/frozen/dynamic
    filename_path = Path(filename).resolve()
    # track if in project folder or in any tracked package directory
    in_project = str(filename_path).startswith(str(PROJECT_ROOT))
    in_package = any(str(filename_path).startswith(str(d)) for d in PACKAGE_DIRS)
    if not (in_project or in_package):
        return trace
    key = f"{frame.f_code.co_name} ({filename_path.name}:{frame.f_code.co_firstlineno})"
    if event == "call":
        frame_stack[frame] = (key, time.perf_counter())
    elif event == "return":
        if frame in frame_stack:
            key, start = frame_stack.pop(frame)
            stats[key]["calls"] += 1
            stats[key]["total_time"] += time.perf_counter() - start
    return trace


def print_stats():
    if not stats:
        print("No functions recorded.")
        return
    sorted_stats = sorted(stats.items(), key=lambda x: x[1]["total_time"], reverse=True)
    col_width = 80
    print(
        "\nFunction profiling results (project + selenium + dependencies, ignoring built-ins/frozen):"
    )
    print(
        f"{'Function':{col_width}s} {'Calls':>5s} {'Total Time':>12s} {'Avg Time':>12s}"
    )
    print("-" * (col_width + 5 + 12 + 12 + 3))
    for key, data in sorted_stats:
        avg = data["total_time"] / data["calls"]
        print(
            f"{key:{col_width}s} {data['calls']:5d} {data['total_time']:12.6f} {avg:12.6f}"
        )


sys.setprofile(trace)
atexit.register(print_stats)