decorators/profile.py (31df010) - testpython.git

profile.py

#!/usr/bin/env python3
"""
A profiler decorator.

Author: Giampaolo Rodola' <g.rodola [AT] gmail [DOT] com>
License: MIT
"""
import cProfile
import tempfile
import pstats
from typing import Union, Optional, Tuple, List, Callable

def profile(sort: Union[Tuple, List, str]='cumulative',
            lines_to_print: Union[int, float]=20,
            strip_dirs: bool=True,
            filename: Optional[str]=None) -> Callable:
    """A decorator which profiles a callable.
    Output goes to stdout if filename is None, otherwise to filename."""
    def outer(fun: Callable) -> Callable:
        def inner(*args, **kwargs):
            prof = cProfile.Profile()
            ret = prof.runcall(fun, *args, **kwargs)
            if filename is None:
                stats = pstats.Stats(prof)
            else:
                statsfile = open(filename, 'a')
                stats = pstats.Stats(prof, stream=statsfile)
            if strip_dirs:
                stats.strip_dirs()
            if isinstance(sort, (tuple, list)):
                stats.sort_stats(*sort)
            else:
                stats.sort_stats(sort)
            stats.print_stats(lines_to_print)
            if filename is not None:
                statsfile.close()

return ret
        return inner

# in case this is defined as "@profile" instead of "@profile()"
    if hasattr(sort, '__call__'):
        fun = sort
        sort = 'cumulative'
        outer = outer(fun)
    return outer

testpython.git

Add a profile decorator.