How To Create Your Own Timing Context Manager In Python

What Is Already Available?¶

Python provides different ways to measure the execution time. For instance, you can use Python’s timeit module [²] 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 [³] 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 [⁴]. 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.

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