Python profiling and timing utils

These are just a few Python snippets that I have been using a lot for timing and profiling functions using only the standard library. You can also find the code here as a GitHub Gist

import contextlib
from functools import wraps
import cProfile
from typing import Optional
import io
import pstats
import time
from pathlib import Path
from functools import wraps


@contextlib.contextmanager
def profile(filename: Path, *args, **kwargs):
    profile = cProfile.Profile(*args, **kwargs)

    profile.enable()
    yield
    profile.disable()

    s = io.StringIO()
    sortby = pstats.SortKey.CUMULATIVE
    ps = pstats.Stats(profile, stream=s).strip_dirs().sort_stats(sortby)
    ps.print_stats()
    with open(filename.with_suffix(".txt"), "w") as f:
        f.write(s.getvalue())

    profile.dump_stats(filename.with_suffix(".prof"))


def profiled(name: Path):
    def decorator(fn):
        @wraps(fn)
        def wrapper(*args, **kwargs):
            with profile(name):
                return fn(*args, **kwargs)

        return wrapper

    return decorator


@contextlib.contextmanager
def timeit(output_file: Optional[Path] = None):
    start = time.perf_counter()
    yield
    end = time.perf_counter()
    if output_file is not None:
        with open(output_file, "a") as f:
            f.write(f"{end - start}\n")
    else:
        print(f"{end - start}")


def timed(output_file: Optional[Path] = None):
    def decorator(fn):
        @wraps(fn)
        def wrapper(*args, **kwargs):
            with timeit(output_file):
                return fn(*args, **kwargs)

        return wrapper

    return decorator

Usage

You can use a context manager or the decorator both for profiling and timing

# context manager
def my_function():
    with profile(Path("/home/ubuntu/profiles/prof")):
        return 1


# decorator
@profiled(Path("/home/ubuntu/profiles/prof"))
def my_function():
    return 1

# context manager
def my_function():
    with timeit(Path("/home/ubuntu/profiles/execution-times.txt")):
        return 1


# decorator
@timed(Path("/home/ubuntu/profiles/execution-times.txt"))
def my_function():
    return 1

When timing functions, it’s always better to run the function multiple times. You can then write a Python script to read all the values in /home/ubuntu/profiles/execution-times.txt (or wherever you have saved the data) and extract some statistics about the function execution time (mean, standard deviation, etc.)

rand[om]

Python profiling and timing utils

Usage