The Magic of Python Context Managers
Tags: python, tutorial, webdev
Published: 2020-10-01 18:15:47


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.

```python
# 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:

```python
# 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: 

```python
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:

```python
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: 

```python
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:

```python
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:

```python
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`:

```python
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:

```python
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`:

```python
# 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:

```python

# 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:

```python
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:

```python
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:

```python
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. &#x1F642;