Resource management is one of those things you need to do in any programming language. Whether you are dealing with locks, files, sessions or database connections - you always have to make sure you close and free up these resources for them operate correctly. Usually, one would do that using try/finally - using the resource in try block and disposing of it in finally block. In Python however, there is a better way - the context management protocol implemented using with statement.

So, in this article we will explore what it is, how it works and most importantly where you can find and how you can implement your own awesome context managers!

What is Context Manager?

Even if you haven't heard of Python's context manager, you already know - based on the intro - that it's replacement for try/finally blocks. It's implemented using with statement commonly used when opening files. Same as with try/finally, this pattern was introduced to guarantee that some operation will be performed at the end of the block, even if exception or program termination occurs.

On surface the context management protocol is just with statement that surrounds block of code. In reality it consists of 2 special (dunder) methods - __enter__ and __exit__ - which facilitate setup and teardown respectively.

When the with statement is encountered in the code, the __enter__ method is triggered and its return value is placed into variable following the as qualifier. After the body of with block executes, the __exit__ method is called to perform teardown - fulfilling the role of finally block.

# Using try/finally
import time

start = time.perf_counter()  # Setup
try:  # Actual body
    time.sleep(3)
finally:  # Teardown
    end = time.perf_counter()
    elapsed = end - start

print(elapsed)

# Using Context Manager
with Timer() as t:
    time.sleep(3)

print(t.elapsed)

The code above shows both the version using try/finally and more elegant version using with statement to implement simple timer. I mentioned above, that __enter__ and __exit__ are needed to implement such context manager, but how would we go about creating them? Let's look at the code of this Timer class:

# Implementation of above context manager
class Timer:
    def __init__(self):
        self._start = None
        self.elapsed = 0.0

    def start(self):
        if self._start is not None:
            raise RuntimeError('Timer already started...')
        self._start = time.perf_counter()

    def stop(self):
        if self._start is None:
            raise RuntimeError('Timer not yet started...')
        end = time.perf_counter()
        self.elapsed += end - self._start
        self._start = None

    def __enter__(self):  # Setup
        self.start()
        return self

    def __exit__(self, *args):  # Teardown
        self.stop()

This code snippet shows Timer class which implements both __enter__ and __exit__ methods. The __enter__ method only starts the timer and returns self which would get assigned in the with ... as some_var. After body of with statement completes, the __exit__ method is invoked with 3 arguments - exception type, exception value and traceback. If everything goes well in body of with statement, those will be all equal to None. If exception gets raised, these are populated with exception data, which we can handle in __exit__ method. In this case we omit exception handling and just stop the timer and calculate elapsed time, storing it in context manager's attribute.

We already saw here both implementation and example usage of with statement, but to have little more visual example of what really happens, let's look at how these special methods get called without the Python's syntax sugar:

manager = Timer()
manager.__enter__()  # Setup
time.sleep(3)  # Body
manager.__exit__(None, None, None)  # Teardown
print(manager.elapsed)

Now that we established what context manager is, how it works and how to implement it, let's look at the benefits of using it - just to have a little more motivation to switch from try/finally to with statements.

First benefit is that whole setup and teardown happens under control of a context manager object. This prevents errors and reduces boilerplate code, which in turn makes APIs safer and easier to use. Another reason to use it is that with blocks highlights the critical section and encourages you to reduce the amount of code in this section, which is also - generally - a good practice. Finally - last but not least - it's a good refactoring tool which factors out common setup and teardown code and moves it into single place - the __enter__ and __exit__ methods.

With that said, I hope I persuaded you to start using context managers instead of try/finally if you didn't use them before. So, let's now see some cool and useful context managers which you should start including in your code!

Making it Simple Using @contextmanager

In the previous section we explored how context manager can be implemented using the __enter__ and __exit__ methods. That's simple enough, but we can make it even simpler using contextlib and more specifically using @contextmanager.

@contextmanager is a decorator that can be used to write self-contained context-management functions. So, instead of creating whole class and implementing __enter__ and __exit__ methods, all we need to do is create single generator:

from contextlib import contextmanager
from time import time, sleep

@contextmanager
def timed(label):
    start = time()  # Setup - __enter__
    print(f"{label}: Start at {start}")
    try:  
        yield  # yield to body of `with` statement
    finally:  # Teardown - __exit__
        end = time()
        print(f"{label}: End at {end} ({end - start} elapsed)")

with timed("Counter"):
    sleep(3)

# Counter: Start at 1599153092.4826472
# Counter: End at 1599153095.4854734 (3.00282621383667 elapsed)

This snippet implements very similar context manager as the Timer class in previous section. This time however, we needed much less code. This little piece of code has 2 parts - everything before yield and everything after yield. The code prior to yield takes the job of __enter__ method and yield itself is the return statement of __enter__ method. Everything after yield is part of __exit__ method.

As you can see above, creating context manager using single function like this requires usage of try/finally, because if exception occurs in body of with statement, it's going to be raised on the line with yield and we will need to handle it in finally block which corresponds to __exit__ method.

As I mentioned already, this can be used for self-contained context managers. It is, however, not suitable for context managers that need to be part of an object, like for example connection or lock.

Even though building context manager using single function forces you to use try/finally and can only be used with simpler use cases, it's still in my opinion elegant and practical option for building leaner context managers.

Real Life Examples

Let's now move on from theory to practical and useful context managers, which you can build yourself.

Logging Context Manager

When the time comes to try to hunt down some bug in your code, you would probably first look in logs to find root cause of the problem. These logs, however, might be set by default to error or warn level which might not be enough for debugging purposes. Changing log level for whole program should be easy, but changing it for specific section of code might be more complicated - this can be solved easily, though, with following context manager:

import logging
from contextlib import contextmanager

@contextmanager
def log(level):
    logger = logging.getLogger()
    current_level = logger.getEffectiveLevel()
    logger.setLevel(level)
    try:
        yield
    finally:
        logger.setLevel(current_level)

def some_function():
    logging.debug("Some debug level information...")
    logging.error('Serious error...')
    logging.warning('Some warning message...')

with log(logging.DEBUG):
    some_function()

# DEBUG:root:Some debug level information...
# ERROR:root:Serious error...
# WARNING:root:Some warning message...

Timeout Context Manager

In the beginning of this article we were playing with timing blocks of code. What we will try here instead is setting timeouts to blocks surrounded by with statement:

import signal
from time import sleep

class timeout:
    def __init__(self, seconds, *, timeout_message=""):
        self.seconds = int(seconds)
        self.timeout_message = timeout_message

    def _timeout_handler(self, signum, frame):
        raise TimeoutError(self.timeout_message)

    def __enter__(self):
        signal.signal(signal.SIGALRM, self._timeout_handler)  # Set handler for SIGALRM
        signal.alarm(self.seconds)  # start countdown for SIGALRM to be raised

    def __exit__(self, exc_type, exc_val, exc_tb):
        signal.alarm(0)  # Cancel SIGALRM if it's scheduled
        return exc_type is TimeoutError  # Suppress TimeoutError


with timeout(3):
    # Some long running task...
    sleep(10)

The code above declares class called timeout for this context manager as this task cannot be done in single function. To be able to implement this kind of timeout we will also need to use signals - more specifically SIGALRM. We first use signal.signal(...) to set handler to SIGALRM, which means that when SIGALRM is raised by kernel our handler function will be called. As for this handler function (_timeout_handler), all it does is raise TimeoutError, which will stop execution in body of with statement if it didn't complete in time. With the handler in place, we need to also start the countdown with specified number of seconds, which is done by signal.alarm(self.seconds).

As for the __exit__ method - if body of context manager manages to complete before time expires, the SIGALRM will be canceled by signal.alarm(0) and program can continue. On the other hand - if signal is raised because of timeout, then _timeout_handler will raise TimeoutError, which will be caught and suppressed by __exit__, body of with statement will be interrupted and rest of the code can carry on executing.

Use What's Already There

Besides the context managers above, there's already bunch of useful ones in standard library or other commonly used libraries like request or sqlite3. So, let's see what we can find in there.

Temporarily Change Decimal Precision

If you're doing lots of mathematical operations and require specific precision, then you might run into situations where you might want to temporarily change precision for decimal numbers:

from decimal import getcontext, Decimal, setcontext, localcontext, Context

# Bad
old_context = getcontext().copy()
getcontext().prec = 40
print(Decimal(22) / Decimal(7))
setcontext(old_context)

# Good
with localcontext(Context(prec=50)):
    print(Decimal(22) / Decimal(7))  # 3.1428571428571428571428571428571428571428571428571

print(Decimal(22) / Decimal(7))      # 3.142857142857142857142857143

Code above demonstrates both option without and with context manager. The second option is clearly shorter and more readable. It also factors-out temporary context which makes it less error prone.

All The Things From contextlib

We already peeked into contextlib when using @contextmanager, but there are more things there which we can use - as a first example let's have a look at redirect_stdout and redirect_stderr:

import sys
from contextlib import redirect_stdout

# Bad
with open("help.txt", "w") as file:
    stdout = sys.stdout
    sys.stdout = file
    try:
        help(int)
    finally:
        sys.stdout = stdout

# Good
with open("help.txt", "w") as file:
    with redirect_stdout(file):
        help(int)

If you have tool or function that by default outputs everything to stdout or stderr, yet you would prefer it to output data somewhere else - e.g. to file - then these 2 context managers might be quite helpful. As in the previous example this greatly improves the code readability and removes unnecessary visual noise.

Another handy one from contextlib is suppress context manager which will suppress any unwanted exceptions and errors:

import os
from contextlib import suppress

try:
    os.remove('file.txt')
except FileNotFoundError:
    pass


with suppress(FileNotFoundError):
    os.remove('file.txt')

It's definitely preferable to handle exceptions properly, but sometimes you just need to get rid of that pesky DeprecationWarning and this context manager will at least make it readable.

Last one from contextlib that I will mention is actually my favourite and it's called closing:

# Bad
try:
    page = urlopen(url)
    ...
finally:
    page.close()

# Good
from contextlib import closing

with closing(urlopen(url)) as page:
    ...

This context manager will close any resource passed to it as argument - in case of the example above - that would be page object. As for what actually happens in the background - the context manager really just forces call to .close() method of the page object the same way as with the try/finally option.

Context Managers for Better Tests

If you want people to ever use, read or maintain test you write you gotta make them readable and easy to understand and mock.patch context manager can help with that:

# Bad
import requests
from unittest import mock
from unittest.mock import Mock

r = Mock()
p = mock.patch('requests.get', return_value=r)
mock_func = p.start()
requests.get(...)
# ... do some asserts
p.stop()

# Good
r = Mock()
with mock.patch('requests.get', return_value=r):
    requests.get(...)
    # ... do some asserts

Using mock.patch with context manager allows you to get rid of unnecessary .start() and .stop() calls and helps you with defining clear scope of this specific mock. Nice thing about this one is that it works both with unittest as well as pytest, even though it's part of standard library (and therefore unittest).

While speaking of pytest, let's show at least one very useful context manager from this library too:

import pytest, os

with pytest.raises(FileNotFoundError, message="Expecting FileNotFoundError"):
    os.remove('file.txt')

This example shows very simple usage of pytest.raises which asserts that code block raises supplied exception. If it doesn't, then test fails. This can be handy for testing code paths that are expected to raise exceptions or otherwise fail.

Persisting Session Across Requests

Moving on from pytest to another great library - requests. Quite often you might need to preserve cookies between HTTP requests, need to keep TCP connection alive or just want to do multiple requests to same host. requests provides nice context manager to help with these challenges - that is - for managing sessions:

import requests

with requests.Session() as session:
    session.request(method=method, url=url, **kwargs)

Apart from solving above stated issues, this context manager, can also help with performance as it will reuse underlying connection and therefore avoid opening new connection for each request/response pair.

Managing SQLite Transactions

Last but not least, there's also context manager for managing SQLite transactions. Apart from making your code cleaner, this context manager also provides ability to rollback changes in case of exception as well as automatic commit if body of with statement completes successfully:

import sqlite3
from contextlib import closing

# Bad
connection = sqlite3.connect(":memory:")
try:
    connection.execute("INSERT INTO employee(firstname, lastname) values (?, ?)", ("John", "Smith",))
except sqlite3.IntegrityError:
    ...

connection.close()

# Good
with closing(sqlite3.connect(":memory:")) as connection:
    with connection:
        connection.execute("INSERT INTO employee(firstname, lastname) values (?, ?)", ("John", "Smith",))

In this example you can also see nice usage of closing context manager which helps dispose of no longer used connection object, which further simplifies this code and makes sure that we don't leave any connections hanging.

Conclusion

One thing I want to highlight is that context managers are not just resource management tool, but rather a features that allows you to extract and factor-out common setup and teardown of any pair of operations, not just common use cases like lock or network connections. It's also one of those great Pythonic features, that you will probably not find in almost any other language. It's clean and elegant, so hopefully this article has shown you the power of context managers and introduced you to a few more ways to use them in your code. 🙂