Profiling Decorators

When things get slow, it's time to find out why! pycallgraph is a great tool to visualize profiling data collected using Python's builtin cProfile module. The same data can also be used in external tools like SnakeViz. Both tend to be a bit tricky to set up, so for your convenience please find some decorator functions for you that do all the work. To produce raw profiling data for external visualization you may use this:

def cprofiled(filepath="profile.prof"):
    """Profile the decorated function and output a .prof file.
 
    Args:
        func (callable): The function to profile.
        filepath (str): The relative profile output file. It can later
            be visualized with tools like snakeviz, gprof2dot etc.
 
    Usage:
        @cprofiled()
        def my_function1(arg1, arg2):
            ...
 
        @cprofiled(filepath="my_func2.prof")
        def my_function2():
            ...
    """
    def decorator(func):
        import cProfile  # noqa
        import pstats  # noqa
        import StringIO  # noqa
 
        def wrapper(*args, **kwargs):
            profile = cProfile.Profile()
            profile.enable()
            result = func(*args, **kwargs)
            profile.disable()
 
            stream = StringIO.StringIO()
            stats = pstats.Stats(profile, stream=stream)
            stats.sort_stats("cumulative")
            stats.dump_stats(filepath)
            return result
        return wrapper
    return decorator

def pycallgraphed(filepath="pycallgraph.png", max_depth=5):
    """Profile the decorated function and output a callgraph image.

    Dependencies:
        - pycallgraph
        - GraphViz's dot executable.

    Args:
        func (callable): The function to profile.
        filepath (str): The relative output image filepath.
        max_depth (int): The maximum call stack level to trace.

    Usage:
        @pycallgraphed()
        def my_function1(arg1, arg2):
            ...

        @pycallgraphed(filepath="output.jpg", max_depth=1)
        def my_function2():
            ...
    """
    def decorator(func):
        from pycallgraph import PyCallGraph  # noqa
        from pycallgraph import Config  # noqa
        from pycallgraph.output import GraphvizOutput  # noqa

        def wrapper(*args, **kwargs):
            config = Config(max_depth=max_depth)
            output = GraphvizOutput(output_file=filepath)
            with PyCallGraph(output=output, config=config):
                return func(*args, **kwargs)
        return wrapper
    return decorator

import time

def timed(func):
    def wrapper(*args, **kwargs):
        start_time = time.time()
        result = func(*args, **kwargs)
        end_time = time.time()
        print("{}() took {} seconds.".format(func.__name__,
                                             end_time - start_time)
        return result
    return wrapper

import time

class Timer(object):
	"""A simple Helper object to measure time in a code block.

	This is simpler to use than the decorators, since we don't need a
	function to wrap etc.

	Example use:

		t = Timer("fetch data")
		...
		t.stop()

	"""
	def __init__(self, name=None):
		self.name = name
		self.start = time.time()

	def stop(self):
		name = u"'{}' ".format(self.name) if self.name is not None else ""
		print(u"{}took {} seconds".format(name, time.time() - self.start))

Home

June 9, 2017, 3 min read

Profiling Decorators