QuantumGhost
7/4/2016 - 3:59 PM

A simple decorator for profiling python function.

A simple decorator for profiling python function.

import cProfile
import functools
import pstats

from cStringIO import StringIO


def profile(func):
    @functools.wraps(func)
    def wrapped(*args, **kwargs):
        print('===== begin profile =====')
        pr = cProfile.Profile()
        pr.enable()
        res = func(*args, **kwargs)
        pr.disable()
        sio = StringIO()
        sortby = 'cumulative'
        ps = pstats.Stats(pr, stream=sio).sort_stats(sortby)
        ps.print_stats()
        print(sio.getvalue())
        print('===== end profile =====')
        return res
    return wrapped