profiler.py

import logging
import time
import typing as T

from cProfile import Profile


log = logging.getLogger(__name__)


def profile(out_filename: str) -> T.Any:
    """ Decorator to profile a function and dump a .profile file.
    Saves the profile file to out_filename

    To prime a function for profiling...

        @profile("/Users/me/Desktop/my-func.profile")
        def my_func():
            ...
            long_running_function()
            ...

    To view the profile stats
        $ pip install tuna
        $ tuna /Users/me/Desktop/my-func.profile
    """
    def wrapper(func: T.Callable) -> T.Callable:
        def profiler(*args: T.Any, **kwargs: T.Any) -> T.Any:
            log.info(f"Profiling: {func.__module__ + '.' + func.__name__}")
            # --- Start profiling ---
            start_ts = time.time()
            profiler_ = Profile()
            return_val = profiler_.runcall(func, *args, **kwargs)
            profiler_.dump_stats(out_filename)
            duration = time.time() - start_ts
            # --- Done profiling ---
            log.info("Finished profiling.")
            log.info(f"{func.__name__}() ran for {duration:.2f}s.")
            log.info(f"Stats dumped at: {out_filename}")
            return return_val
        return profiler
    return wrapper