Article Thumbnail

How To Create Your Own Timing Context Manager In Python

Implementing it the easy way

Florian Dahlitz
5 min
April 26, 2019

Introduction

Updated May 10, 2020: Adjust language and structure and make sure the code is compatible with CPython 3.8.2.

After publishing my last article, which was about list.sort() and sorted(list), I was asked, why I was using the boxx package [1] instead of built-in functionalities to measure the execution time of certain pieces of code. I responded that it is only personal preference and that you can simply create your own context manager measuring the execution time of code pieces.

In this article I will show you how to create your own timing context manager. Furthermore, different ways to accomplish that are covered.

Time Measurement

What Is Already Available?

Python provides different ways to measure the execution time. For instance, you can use Python’s timeit module [2] to measure the execution time of code pieces.

>>> import timeit
>>> timeit.timeit('"-".join(str(n) for n in range(100))', number=10000)
0.3018611848820001

However, the timeit.timeit() function only accepts strings. This can be quite annoying if you want to measure larger functions. The following example is from the official Python documentation [3] and shows you how you can run and measure functions using the timeit module.

def test():
    """Stupid test function"""
    L = [i for i in range(100)]


if __name__ == '__main__':
    import timeit
    print(timeit.timeit("test()", setup="from __main__ import test"))

Although it works, it does not look really pythonic.

Another way to measure the execution time is to make use of Python’s cProfile module [4]. However, this is not recommended! It is meant to be used to get a sense of how long certain code pieces need to be executed and not to measure the exact execution time. In fact it is not really precise. You can use it via:

$ python -m cProfile <file_name.py>

Now that we have seen a recommended but unpythonic, and a discredited and unpythonic way, how can we implement our own recommended and pythonic solution?

The Idea

The idea is quite simple: We take the time at the beginning of the execution and subtract it from the time at the end of the execution. Fortunately, Python has a built-in module we can use: time.

>>> import time
>>> start = time.time()
>>> # do some stuff
>>> end = time.time()
>>> print(f"Elapsed Time: {end - start}")

Great! However, adding one line of code before and two lines after a code piece we want to measure is an overhead I don’t want to have. So let’s create a context manager for that.

Creating a Context Manager

There exist different ways to create a context manager in Python. We will have a look at two ways to accomplish that: A class-based and a generator-based solution.

Class-Based Context Manager

To create a class-based context manager, the dunder methods __enter__ and __exit__ need to be implemented. The first one is called when entering a context (manager), the latter is called when leaving the context.

With this knowledge we can create a Timer class implementing both methods. When entering the context, we want to get the current time and save it to a start instance variable. If we leave the context, we want to get the current time and subtract the start time from it. The result is printed.

To customize the output, we let the user specify a description, which is printed before the elapsed time. The following snippet shows you the ready-to-use class.

# class_based_context_manager.py

from time import time


class Timer(object):
    def __init__(self, description):
        self.description = description

    def __enter__(self):
        self.start = time()

    def __exit__(self, type, value, traceback):
        self.end = time()
        print(f"{self.description}: {self.end - self.start}")


with Timer("List Comprehension Example"):
    s = [x for x in range(10_000_000)]

Executing the file results in something similar to:

$ python class_based_context_manager.py
List Comprehension Example: 0.3361091613769531

Generator-Based Context Manager

The generator-based approach is a bit more straightforward. Basically, we create a generator function containing the program flow (taking start and end time as well as printing the elapsed time). Then, we utilize the contextlib's @contextmanager decorator to turn the generator function into a proper context manager.

# generator_based_context_manager.py

from contextlib import contextmanager
from time import time


@contextmanager
def timing(description: str) -> None:
    start = time()
    yield
    ellapsed_time = time() - start

    print(f"{description}: {ellapsed_time}")


with timing("List Comprehension Example"):
    s = [x for x in range(10_000_000)]

Let’s describe it a bit less formal: What happens is that entering the context (timing) results in taking the current time and playing the ball back to the code inside the context using yield . If the code inside of the with-block is executed, we jump back to the point right after the yield keyword. Now, we calculate the elapsed time and print it. We are now leaving the context.

Executing the example at hand results in something similar to:

$ python generator_based_context_manager.py
List Comprehension Example: 0.3237767219543457

Summary

In this article you learned how to create your own timing context manager. After having a look at the basic concept, you implemented the timing context manager in two ways: Class-based and generator-based. The resulting class and generator function are ready-to-use.

I hope you enjoyed reading the article. Make sure to share it with your friends and colleagues. If you have not already, consider following me on Twitter, where I am @DahlitzF or subscribing to my newsletter so you won’t miss a future article. Stay curious and keep coding!

References