Profiling context management

profile_ctx.py

import cProfile
import contextlib
import io
import pstats
import sys
import timeit


@contextlib.contextmanager
def prof(*restrictions, stdout=True, dump=None, sortby='cumulative'):
    """Profile code running in this context.
    
    Arguments:
      restrictions: Passed down to https://docs.python.org/3.6/library/profile.html#pstats.Stats.print_stats
      stdout: write profile stats to stdout
      dump: path to dump  pstats (optional)
      sortby: Sorting criteria: https://docs.python.org/3.6/library/profile.html#pstats.Stats.sort_stats

    Usage:
      >>> # this will print top 10 calls sorted by 'cumtime' to stdout.
      >>> with prof(10, sortby='cumulative'):
      >>>     do_stuff()
      >>>     do_other_stuff()
    """
    pr = cProfile.Profile()
    pr.enable()
    yield
    pr.disable()
    if stdout:
        s = io.StringIO()
        ps = pstats.Stats(pr, stream=s).sort_stats(sortby)
        ps.print_stats(*restrictions)
        print(s.getvalue())
    if dump:
        pr.dump_stats(dump)


@contextlib.contextmanager
def time_it(name='default_timer', stream=sys.stdout):
    start = timeit.default_timer()
    yield
    stream.write('time_it ({}) - {}\n'.format(name, timeit.default_timer() - start))