37
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!
with
statement allows us to reduce code duplicationIf 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.
Why is that bad? Leaving files or stateful resources open unnecessarily is bad for the following reasons (source):
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!')
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.
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.
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()
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:
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
Here’s a better example where we can use context manager with
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.If you have made it this far, awesome! In summary, here's what we've learned:
with
statement helps to encapsulate the standard use of try
, finally
, else
when it comes to exception handling.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!
37