Profiling stuff

profiling.py

"""Profiling code."""

import cProfile
import io
import pstats
import contextlib


def start():
    """Start and return profiler."""
    profiler = cProfile.Profile()
    profiler.enable()
    return profiler


def finish(profiler):
    """Stop the profiler and print out stats."""
    profiler.disable()
    out_stream = io.StringIO()
    profile_stats = pstats.Stats(
        profiler, stream=out_stream).sort_stats('cumulative')
    profile_stats.print_stats(30)
    print(out_stream.getvalue())


@contextlib.contextmanager
def profiled(enabled):
    """Context manager to profile within a given context."""
    profiler = None
    if enabled:
        profiler = start()
    try:
        yield
    finally:
        if profiler is not None:
            finish(profiler)

This would be great if it included the option of dumping the stats to a file. You can do that in finish by calling profile_stats.dump_stats.

Then users can load the stats file at the console using pstats.Stats and have the option of sorting the data in different ways. They can also pass the file to SnakeViz to visualize the call tree and execution time.