decorators/profile.py (010f8dd) - 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 pstats
import functools
from typing import Union, Optional, Tuple, List, Callable

def profile(sort: Union[Tuple, List, str]='cumulative',
            lines_to_print: float=20,
            strip_dirs: bool=True,
            filename: Optional[str]=None) -> Callable:
    """A decorator which profiles a callable.
    If lines_to_print is an int, it's the number of lines. If it's a float
    then it's the percentage of total lines. (Eg., .2 = 20%)
    Output goes to stdout if filename is None, otherwise to filename."""
    def outer(fun: Callable) -> Callable:
        @functools.wraps(fun)
        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)  # type: ignore
    return outer

testpython.git

Cleanup for basedpyright