Profile method decorator in Python showing a graphical (.svg) representation (as well as a text-based representation) as the output.

profiling.py

import pstats
import sys
import subprocess
import os
try:
    import cProfile as profile
except ImportError:
    import profile


def profile_method(filename, rows=50, sort=(('cumul',), ('time',)), show_graph=False):
    '''
    Decorator to profile the decorated function or method.
    
    To use:
    
    @profile_method('out.prof', show_graph=True)
    def my_func_to_profile():
        ...
        
    my_func_to_profile()


    Depends on:  

    - Graphviz to generate the graph: http://www.graphviz.org/ must be installed with the GRAPHVIZ_DOT environment variable pointing to dot.exe. 

    - gprof2dot: http://gprof2dot.jrfonseca.googlecode.com/git/gprof2dot.py must be in the PYTHONPATH). 

    - desktop: optional: https://pypi.python.org/pypi/desktop if available will open the generated .svg automatically (a viewer must be properly registered)
    

    :param int rows: 
        How many rows of data to print.
    
    :param tuple sort: 
        How to sort the stats of the printed data. 
        
        Options available from: 
        stats = pstats.Stats(prof)
        print(stats.get_sort_arg_defs().keys())
        
    :param bool show_graph:
        Whether a graph should be generated. Note: the computer should have an .svg viewer
        associated to the extension so that the file is properly opened.
    '''

    def wrapper(method):
        def inner(*args, **kwargs):
            prof = profile.Profile()
            result = prof.runcall(method, *args, **kwargs)

            prof.dump_stats(filename)

            if show_graph:
                _show_graph(filename)

            # Show text output regardless of showing graph.
            tup_sort = sort
            s = tup_sort[0]
            if isinstance(s, str):
                tup_sort = [tup_sort]

            stats = pstats.Stats(prof)
            for s in tup_sort:
                stats.strip_dirs().sort_stats(*s).print_stats(int(rows))

            return result
        return inner
    return wrapper


def _show_graph(filename):
    '''
    Creates an .svg from the profile generated file and opens it (a proper association to .svg
    files must be already defined in the machine).
    
    @param str filename:
        This is the file generated from profile_method.
    '''
    import gprof2dot
    initial = sys.argv[:]
    output_filename = filename + '.dot'
    sys.argv = ['', '-o', output_filename, '-f', 'pstats', filename]
    try:
        gprof2dot.Main().main()
    finally:
        sys.argv = initial

    try:
        dot = os.environ['GRAPHVIZ_DOT']
    except KeyError:
        raise AssertionError('The GRAPHVIZ_DOT environment variable must be defined to show graph.')

    assert os.path.exists(dot), 'Expected: %s to exist and point to dot.exe' % dot
    subprocess.call([dot, '-Tsvg', '-O', output_filename])

    print('Opening svg created at:', os.path.realpath((output_filename + '.svg')))
    try:
        import desktop
    except:
        pass
    else:
        desktop.open(output_filename + '.svg')