profiling-utils.py

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.py

# 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