guohanghui/graph-theory

1import cProfile2import io3import pstats4 5 6def profileit(func):7    """ decorator for profiling function.8    usage:9    >>> def this(var1, var2):10            # do something11 12    >>> new_this = profileit(this)13    >>> calls, cprofile_text = new_this(var1=1,var2=2)14    """15    def wrapper(*args, **kwargs):16        prof = cProfile.Profile()17        retval = prof.runcall(func, *args, **kwargs)18        s = io.StringIO()19        sortby = 'cumulative'20        ps = pstats.Stats(prof, stream=s).sort_stats(sortby)21        ps.print_stats()22        text = s.getvalue()23        calls = int(text.split('\n')[0].lstrip().split(" ")[0])24        return calls, text25    return wrapper26