DEV Community

Jerry Ng
Jerry Ng

Posted on • Edited on • Originally published at jerrynsh.com

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!')
Enter fullscreen mode Exit fullscreen mode

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 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()
Enter fullscreen mode Exit fullscreen mode

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 fullscreen mode Exit fullscreen mode

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.
Enter fullscreen mode Exit fullscreen mode

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()
Enter fullscreen mode Exit fullscreen mode

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()
Enter fullscreen mode Exit fullscreen mode

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!
Enter fullscreen mode Exit fullscreen mode

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!


This article was originally published at jerrynsh.com

Top comments (4)

Collapse
 
mccurcio profile image
Matt Curcio

Hi Jerry,
Helpful article. I learned something new. I have seen the with and the common file handle open approach but did not realize the differences could effect your program later on.

May I make one suggestion. I think your conclusion is very important. BUT, I am reminded of the 'old' advice about writing papers.

  1. Tell them what you are going to say.
  2. Tell them.
  3. Tell them what you said.

If believe the article might have had more impact on if the conclusion was ALSO at the beginning. This could help frame the text into 4 bullets. The intro need not say it the same way as the conclusion but I find it reinforces the points better for me.
Just my 2 cents ;)

Collapse
 
jerrynsh profile image
Jerry Ng

Hey Matt, I'm glad you learn something new!
I love your feedback and I'll do that in my future articles! I really appreciate it!

Collapse
 
incrementis profile image
Akin C.

Hello Jerry Ng,

thanks for your article.
I am currently learning python for my upcoming studies.
I've read your code examples and coded them at the same time.
I noticed that the first snippet of code only worked when
'data = f.read ()' is removed.

I understand that this article is not intended for beginners, but it still has increased my knowledge of python a bit 🤓!

Collapse
 
jerrynsh profile image
Jerry Ng

Hey Akin! Good catch!

It was using w only which stands for Write-only. Instead, we should use w+ if we want to allow Read and Write operations.

I've updated it to be:

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