Clean Resource Management in Python using Context Managers
Tags: python
Published: 2019-12-20 06:15:56

Managing resources is a vital aspect of software. Closing open files, freeing dynamic memory, terminating database sessions. Although it's tedious, we know it's necessary.

But wouldn't it be nice if we could delegate resource allocation and release, avoiding those cumbersome `try...finally` blocks?

Well, Python has our backs thanks to the `with` statement and context managers!

### The `with` statement

Many of you will be familiar with the 'pythonic' way of opening a file, working with it and then making sure it gets closed; which would be something like:

``` python
with open('my-file.txt') as my_file:
    # Do your thing
```

That way you open the file, work with it and know it will be automatically closed once you are finished; even if your code raises an error. That code is equivalent to:

``` python
my_file = open('my-file.txt')
try:
    # Do your thing
finally:
    my_file.close()
```

However, as you can see; the first solution is much cleaner and easier to understand. That version is pure business logic, you open the file and do something with it. It doesn't include the fact that the file needs to be closed or that an error might occur while working with the file like the second version, it's all done behind the scenes.

Another instance when you can use the `with` statement is while using lock object from the threading module.

``` python
lock = Lock()

# Why write this?
lock.acquire()
try:
    ...
finally:
    lock.release()

# When you can write this
with lock:
    ...
```

Now, the question is, can you manage your custom resources using the `with` statement?

Yes! You can!

### Context managers

A context manager is any object that specifies two operations that should be executed one after the other, with some arbitrary logic between them. It does so through the methods `__enter__()` and `__exit__()`. Note that we are talking about any two operations, not necessarily acquisition and release of a resource.

These context managers can be used inside the `with` statement as follows:

``` python
with some_context_manager [as foo]:
    # Your code
```

When Python executes that piece of code:

1. `some_context_manager.__enter__()` is called. The method is passed no arguments and, if the `as` clause was used, it's return value is bound to the name `foo` inside the inner block.

2. Your code get executed, having access to `foo` if specified.

3. `some_context_manager.__exit__()` is called. It is passed 3 arguments for handling any exception raised by your code. If no exception was raised, all three are `None`. These are:
    * `exc_type`: the type of the raised exception.
    * `exc`: the exception itself.
    * `exc_traceback`: the exception's traceback.

    If `__exit__()` wishes to supress the exception raised by the inner code block, preventing it from being propagated, it should return `True`.

### Writing your own context manager

As we described before, a context manager is just and object that responds to the `__enter__()` and `__exit__()` methods. We can write a simple manager like this:

``` python
class MyContextManager:

    def __enter__(self):
        print('Acquiring a resource')
        self.resource = acquire()
        return self.resource

    def __exit__(self, exc_type, exc, exc_traceback):
        if exc_type is not None:
            print('Oops, this error was raised:', exc)
        print('Releasing the resource')
        self.resource.release()
```

Let's try it out.

``` python
>>> with MyContextManager():
...   print('Using the resource')

# Output
Acquiring a resource
Using the resource
Releasing the resource
```
Now, let us see what happens if we raise an exception inside the `with` statement:

``` python
>>> with MyContextManager():
...   raise ValueError('Something went wrong')

# Output
Acquiring a resource
Oops, this error was raised: Something went wrong
Releasing the resource

Traceback (most recent call last):
  File "<stdin>", line 2, in <module>
ValueError: Something went wrong
```

You can see that, even though we raised an error inside the code block, the __exit__() method was still executed and the resource released before propagating the exception outside of the `with` statement.

What if we wanted to supress the error? In that case, we could modify `__exit__()` and return `True` from in:

``` python
def __exit__(self, exc_type, exc, exc_traceback):
    if exc_type is not None:
        print('Oops, this error was raised:', exc)
    print('Releasing the resource')
    self.resource.release()
    return True
```

And now:

``` python
>>> with MyContextManager():
...   raise ValueError('Something went wrong')

# Output
Acquiring a resource
Oops, this error was raised: Something went wrong
Releasing the resource
```

The exception was raised and `__exit__()` might do something to handle it, but it was not propagated.

### Writing a context manager as a function

Python's built in module [`contextlib`](https://docs.python.org/3/library/contextlib.html) provides utilities for working with context managers. It includes many useful things, such as common context managers already implemented, but in this section we'll focus on the `contextmanager` decorator.

This decorator allows to write a context manager as a function which returns a generator that yields one value. For example, we could write `MyContextManager` like this:

``` python
from contextlib import contextmanager

@contextmanager
def some_resource():
    print('Acquiring a resource')
    resource = acquire()
    try:
        yield resource
    except Exception as exc:
        print('Oops, this error was raised:', exc)
        raise
    finally:
        print('Releasing the resource')
        resource.release()
```

In this case, when you run:

``` python
with some_resource() [as foo]:
    # Your code
```

1. `some_resource()` is executed up until the `yield` statement

2. Your code is executed, having the resource bound to the name `foo` if specified.

3. The rest of `some_resource()` is executed.

I personally prefer this method for writing simple managers. For more complex ones, is generally better to use a class.

### A real world example

I was recently on an automation Python script that ran on an Amazon EC2 instance and, at one point, it needs to execute an `awscli` command using a different AWS account than that of the virtual machine. One of the ways of temporarily using a different account is getting that account's credentials and storing them in two environment variables, `AWS_ACCESS_KEY_ID` and `AWS_SECRET_ACCESS_KEY`.

So; I needed to set those environment variables, run the cli command and unset the variables so I can continue using my original account. Sounds perfect for a context manager!

``` python
from os import environ, unsetenv
from contextlib import contextmanager

@contextmanager
def aws_credentials(access_key: str, secret_key: str):
    environ['AWS_ACCESS_KEY_ID'] = access_key
    environ['AWS_SECRET_ACCESS_KEY'] = secret_key
    try:
        yield
    finally:
        unsetenv('AWS_ACCESS_KEY_ID')
        unsetenv('AWS_SECRET_ACCESS_KEY')
```

**Note that using `yield` by itself is equivalent to writing `yield None`.**

As you can see; that context manager sets the environment credentials, yields so that the code inside the `with` statement is executed and then unsets the credentials. It can be used as follows:

``` python
access_key, secret_key = securely_get_credentials()
with aws_credentials(access_key, secret_key):
    # Run cli command
```

### Conclusion

Context managers are a very simple tools that allows you to write code that is a lot cleaner and more concise. And now that you know them, you'll be surprised by how often they show up!

This has been my first post on dev.to, I hope you liked it. If you've got any feedback, feel free to let me know :smile:.

Cheers!