Clean Resource Management in Python using Context Managers

Pablo Paglilla — Fri, 20 Dec 2019 06:15:56 +0000

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:

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:

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.

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:

with some_context_manager [as foo]:
    # Your code

When Python executes that piece of code:

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.
Your code get executed, having access to foo if specified.
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:

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.

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

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

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:

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

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:

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

some_resource() is executed up until the yield statement
Your code is executed, having the resource bound to the name foo if specified.
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!

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:

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 😄.

Cheers!

DEV Community: Pablo Paglilla