All You Need To Know About Context Managers In Python

Regarded as an obscure feature by some, the with statement in Python is often only used when dealing with file operations. Below is an example to read and write a file with Python's built-in open function.

with open('hello.txt', 'w+') as f:
    data = f.read()
    f.write('Hello Python!')

For the most part, all people knew was that using the with statement is the preferred way to manage files rather than using the close method.

The fact is, most people couldn't be bothered enough to peek behind the scene of what is happening behind the scenes. Here's a tip, the underlying protocol is known as a context manager.

To me, it was all magic. Honestly, who would without answering why should we ever use it?

In this article, I am going to share with you why we should use context manager along with some practical use cases involving database connections. You could probably refactor some part of your codebase to use context manager with the with statement as well.

Let's dive into it!

TL;DR

Avoid leaving any files or database connections open as they are limited

Context manager allows us to better manage these resources by telling an object what to do when created or destroyed

The use of with statement allows us to reduce code duplication

Whys

Resource management

If I were to summarize it, it would be just two words: resource management.

When building any applications, it's common for us to use resources like file operations and database connections. Here's a key takeaway, these resources are limited.

Oftentimes, we would need to "release" these resources after using them. As an example, whenever we open a file from our filesystem, we need to explicitly close the file when we are done using it.

Don't leave of leaving files or resources open

Why is that bad? Leaving files or stateful resources open unnecessarily is bad for the following reasons (source):

They may consume limited system resources, such as file descriptors. Code that deals with many such objects may exhaust those resources unnecessarily if they’re not returned to the system promptly after use.

Holding files open may prevent other actions such as moving or deleting them, or unmounting a filesystem.

Files and sockets that are shared throughout a program may inadvertently be read from or written to after logically being closed. If they are actually closed, attempts to read or write from them will raise exceptions, making the problem known sooner.

The "with" statement

So, what is the with statement or context manager good for you ask?

Sure, there is nothing wrong with calling session.close() every time we are done with our database transaction in sqlalchemy. Nor is there anything wrong with having to call the built-in close method every single time we are done reading and writing a file. For that matter, here’s an example of one of these methods:

# Poor Example
# ------------
f = open('hello.txt', 'w')
f.write('Hello Python!')
f.close()

As you can already tell, the example given above is quite verbose. Now, imagine doing it in every single part of your codebase (gross, I know).

Besides, there’s a good chance that a poor, tired developer might just forget to close the file (or a database connection) after using it.

Hence, opening a file using the with statement is generally recommended. Using with statements help you to write more expressive code while avoiding resource leaks.

# Good Example
# ------------
with open('hello.txt', 'w') as f:
    f.write('Hello Python!')

Enter Context Managers

Resource management can be achieved by using context managers in Python. In essence, context managers help to facilitate proper handling of resources, providing users mechanism for setup and teardown of resources easily.

To reiterate in layman terms, context managers allow you to control what to do when objects are created or destroyed.

There are several ways to create a reusable context manager in Python. In the next section, I am going to run through several examples of how you can create context managers in Python.

For the first two examples, let's create a simple custom context manager to replace the built-in open function in Python.

Do note that In practice, we should always use any built-in methods or context manager that is provided by Python.

1. Class based

The classic example would be creating a Python class for your own context manager. By default, every context manager class must contain these three Dunder methods:

__init__

__enter__

__exit__

These methods will be executed sequentially as shown above. Please refer to the comments in the code example below for a more detailed explanation.

Note that the code below can only serve as an example and should NOT be used to replace the use of the built-in open function.

class CustomFileHandlerContextManager:
    """
    A custom context manager used for handling file operations
    """

    def __init__(self, filename, mode):
        print('__init__ method is called.')

        self.filename = filename
        self.mode = mode
        self.file = None

    def __enter__(self):
        print('__enter__ method is called.')

        self.file = open(self.filename, self.mode)
        return self.file

    def __exit__(self, exc_type, exc_value, exc_traceback):
        print('__exit__ method is called.')

        self.file.close()  # NOTE: So that we can use `CustomFileHandlerContextManager('hello.txt', 'w') as f`


def main():
    with CustomFileHandlerContextManager('hello.txt', 'w') as f:  # __init__ and __enter__ is called
        f.write('Hello! I am not Tom!')
        print('Do something else in the statement body.')

        # __exit__ is called upon exception or end of the `with` statement

    assert f.closed is True  # Proof that the file is closed :)


if __name__ == '__main__':
    main()

# Output:

# __init__ method is called.
# __enter__ method is called.
# Do something else in the statement body.
# __exit__ method is called.

2. Generator based

Another popular alternative to writing a context manager is to use the built-in contextlib library in Python. It is my personal preferred way of creating a custom context manager.

As an overview, contextlib provides us a set of utilities for common operations involving the with statements.

With contextlib, we can omit writing a Python class along with the required Dunder methods for our custom context managers.

import contextlib

@contextlib.contextmanager
def custom_file_handler(file_name, file_mode):
    file = open(file_name, file_mode)

    yield file # NOTE: So that we can use `custom_file_handler('hello.txt', 'w') as f`

    file.close() # Anything after yield will act is if it's in the __exit__

def main():
    with custom_file_handler('test.txt', 'w') as f:
        f.write('Hello, I am Jerry! This is a generator example.')

        print('Do something else in the statement body.')

    assert f.closed is True # Proof that the file is closed :)

if __name__ == '__main__':
    main()

3. Use built-in context managers

Generally speaking, we should avoid re-inventing the wheel. We should always opt for using any available built-in context managers if there were made available.
For instance, if you're working with SQLAlchemy, the library already provides a good way to manage sessions. When working with SQLAlchemy ORM, it's a common pattern for us to do this:

from sqlalchemy import create_engine
from sqlalchemy.exc import ProgrammingError
from sqlalchemy.orm import sessionmaker


engine = create_engine('postgresql://jerry:nsh@localhost/')
Session = sessionmaker(engine)
session = Session()

try:
    session.add(some_object)
    session.add(some_other_object)

except ProgrammingError as e:
    logger.exception(e)
    session.rollback()
    raise

else:
    session.commit()

Instead of having to call session.rollback() and session.commit() every single time across numerous functions, we can instead use the built-in session as a context manager.
Here’s a better example where we can use context manager with sessionmaker:

from sqlalchemy import create_engine
from sqlalchemy.exc import ProgrammingError
from sqlalchemy.orm import sessionmaker


engine = create_engine('postgresql://jerry:nsh@localhost/')
db_session = sessionmaker(engine)

# We can now construct a session which includes begin(), commit(), rollback() all at once
with db_session.begin() as session:
    session.add(some_object)
    session.add(some_other_object)
# Commits the transaction, closes the session auto-magically! Cool!

Using the with statement here makes our code look much, much cleaner.

Closing Thoughts

If you have made it this far, awesome! In summary, here's what we've learned:

We should not leave any stateful resources (files, database connections, sockets) open unnecessarily as they are limited.

Python's context manager allows us to better manage these resources by telling an object what to do when created or destroyed.

with statement helps to encapsulate the standard use of try, finally, else when it comes to exception handling.

The use of a context manager allows us to reduce code duplication.

Using Python's context manager with the with statement is a great choice if your code has to deal with the opening and closing of a file or database connection.

Personally, my favorite part of using context managers is that it allows us to simplify some common resource management patterns. Context managers abstract their own functionality, thus allowing them to be refactored out and reused repeatedly.

That is all! Thank you for reading!