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

Using the power of graphviz, we can visualize the results in the same step:

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

And for when you just want to know execution time:

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

For situations where you don’t have a function to decorate and don’t want to refactor your code in that way this one is also useful:

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))

Advertisements

Leave a Reply

Fill in your details below or click an icon to log in:

WordPress.com Logo

You are commenting using your WordPress.com account. Log Out / Change )

Twitter picture

You are commenting using your Twitter account. Log Out / Change )

Facebook photo

You are commenting using your Facebook account. Log Out / Change )

Google+ photo

You are commenting using your Google+ account. Log Out / Change )

Connecting to %s