A profiling context manager for python

gistfile1.txt

import arrow
import cProfile
import contextlib
import os


@contextlib.contextmanager
def profiled(dump_file_prefix='my-profile.prof', dump_file_dir='prof', logger=None):
    """
    Profiles a section of code for later analysis. Use as a context manager
    as in the following example:

    .. code-block:: python

        with profiled(dump_file_prefix='my-profile', dump_file_dir='prof'):
            do_something_cool()

    After the function is called, a cProfile dump file is available in the
    'prof' directory with a filename that starts with 'my-profile'. The output
    can be explored with the awesome tool `snakeviz
    <https://jiffyclub.github.io/snakeviz/>`_.
    """
    if not os.path.exists(dump_file_dir):
        os.mkdir(dump_file_dir)
    dump_filename = "{}/{}.{}.prof".format(
        dump_file_dir,
        dump_file_prefix,
        arrow.now().format('YYYYMMDD.HHmmss.SSS'),
    )
    pr = cProfile.Profile()
    pr.enable()
    yield
    pr.disable()
    pr.dump_stats(dump_filename)
    if logger is not None:
        logger.info("Profile written to {}", dump_filename)