2016-03-08

A Python Timer as Context Manager

The timing problem

How often have you had the need to time a piece of code? I very often had something like

import time

start_time = time.time()

long_running_function(...)

duration = time.time() - start_time
print("Duration", duration)

It is kind of annoying to have this repeatedly. Luckily, Python wasn't Python if there wasn't an easier way around this. One possible solution is to use a context manager. A context manager is the thing you put into Python's with-statement. Let's see how the code would look like from a user perspective before checking out the implementation of such a timer context manager. What I would like to write is something like

>>> with Timer() as timer:
>>>    long_running_function()
<Timer duration=30>
>>> timer.duration
30

A context manager is any object implementing the special methods __enter__ and __exit__. The __enter__ method is called on entering the context and the __exit__ method when the context is left. Here is one possible implementation:

import time


class Timer:
    def __init__(self, print_at_exit=True):
        self.print_at_exit = print_at_exit

    def __enter__(self):
        self.start_time = time.time()
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        self.end_time = time.time()
        self.exc_type = exc_type
        if self.print_at_exit:
            print(self)

    def __repr__(self):
        return "<{} duration={}{}>".format(self.__class__.__name__, self.duration, self.exception())

    def exception(self):
        return " exception={}".format(self.exc_type.__name__) if self.exc_type else ""

    @property
    def duration(self):
        try:
            return self.end_time - self.start_time
        except AttributeError:
            return -1

At object initialization we can additionally specify if we want a print of the elapsed time on exiting the context. The __enter__ method returns the thing that is bound to whatever follows the "as". In our case this is the timer variable. The __exit__ method gets additional information about possible exception which might have occurred in the context. We don't handle these but display the exception type in Timer's __repr__ method. Lastly, the duration is calculated on the fly as a property.

Context managers in general

Context managers complement functions in the sense that functions are surrounded by code, whereas context managers surround code themselves.