cProfiling.py function decorator

cProfiling.py

#!/usr/bin/env python3
"""cProfile decorator to get printed and optionally dumped statistics."""

import cProfile, pstats
from functools import wraps

def profile(output_file=None, sort_by='cumulative', lines=None,
            dump_stats=False, reverse=False, strip_dirs=False):
    """A time profiler decorator.

    Inspired by and modified from the profile decorator of Giampaolo Rodola:
        http://code.activestate.com/recipes/577817-profile-decorator/

    Args:
        output_file: str or None. Default is None
            Path of the output file. If only name of the file is given, it's
            saved in the current directory.
            If it's None, the name of the decorated function and sortby is used.
        sort_by: str or SortKey enum or tuple/list of str/SortKey enum
            Sorting criteria for the Stats object.
            For a list of valid string and SortKey refer to:
            https://docs.python.org/3/library/profile.html#pstats.Stats.sort_stats
        lines: int or decimal or None
            Number of lines to print. Default (None) is for all the lines.
            This is useful in reducing the size of the printout, especially
            that sorting by 'cumulative', the time consuming operations
            are printed toward the top of the file.
            For 'decimal', e.g.: .1 can be used for 10% of the total lines.
        dump_stats: bool, default: False
            Dump the Stats object data to file named as output_file + '.dump'.
        reverse: bool, default: False
            Reverse sorting from descending to ascending order.
        strip_dirs: bool, default: False
            Whether to remove the leading path info from file names.
            This is also useful in reducing the size of the printout

    Returns:
        Profile of the decorated function

    Usage Example:
        @profile(output_file='profile_func-asc-cumtime.stats',
                 sort_by=('cumulative', 'time'),
                 lines=10, strip_dirs=True,
                 reverse=True)
        def the_func():
            print('Life is good.')
    """
    def inner(func):
        @wraps(func)
        def wrapper(*args, **kwargs):
            _outsorts = '-'.join(sort_by) if isinstance(sort_by, (tuple, list)) else sort_by
            _output_file = output_file or func.__name__ + '.profile.' + _outsorts + '.stats'

            # Get the profile data and run the 'func'
            pr = cProfile.Profile()
            pr.enable()
            retval = func(*args, **kwargs)
            pr.disable()

            if dump_stats:
                pr.dump_stats(_output_file + '.dump')

            with open(_output_file, 'w') as f:
                ps = pstats.Stats(pr, stream=f)
                if strip_dirs:
                    ps.strip_dirs()
                if isinstance(sort_by, (tuple, list)):
                    ps.sort_stats(*sort_by)
                else:
                    ps.sort_stats(sort_by)
                if reverse:
                    ps.reverse_order()
                ps.print_stats(lines)
            return retval

        return wrapper

    return inner