Python context managers are a nifty resource management tool that provides a simple syntax for a powerful construct. When working with threads, database connections, files, or any limited resources, it’s essential to release them after use. Not releasing these resources leads to resource leakage and can cause the program to act erratically or even crash. Generally, the allocating and releasing of resources is achieved by using try/finally, utilizing the resource in the try block and releasing it in the finally block. Python, however, has a better, more convenient way. And even if this is the first time you’ve come across the term “context managers”, you’re probably familiar ‘with’ it. This article will explore what context managers are, how they work, how to implement your own and then some use-cases.
What is a Context Manager?
Python provides the with statement that enables the developer to implement the concept of a runtime context. This can be broken down and understood using the analogy of a local variable inside a function. The local variable is created when the function is called and is, ideally, automatically deleted when the function has finished execution. In this case, we are releasing memory when we’re done with it. Similarly, in the case of context managers, the resource is locked or marked as ‘in use’ when the interpreter enters the context manager block and released as soon as it’s exited.
Let’s see the trusty with statement in action and compare it with its more verbose try/finally alternative:
with open('the-tales-of-beadle-the-bard.txt', 'w') as f:
f.write("If Harry Potter speaks Parseltongue, can he code in Python?")
is equivalent to:
try:
f = open('the-tales-of-beadle-the-bard.txt', 'w')
f.write("If Harry Potter speaks Parseltongue, can he code in Python?")
finally:
f.close()
How to Write Custom Python Context Managers
There are two ways of implementing custom context managers in Python:
- implementing a context manager class using the
__enter__and__exit__dunder methods - creating a generator using the
contextlib.contextmanagerdecorator
As a Class
Context managers can be implemented as classes using the two dunder methods __enter__ and __exit__. The with statement calls the __enter__ method, and the value returned by this method is assigned to the identifier in as clause of the statement. After the __enter__ method is called, the code inside the with block is executed, and then the __exit__ method is called with three arguments – exception type, exception value and traceback. All three arguments are None if the code inside the with block doesn’t raise any exceptions. If there’s is an exception it can be handled in the __exit__ method using the three arguments. Let’s illustrate this by creating a timed file handling context manager.
from time import perf_counter, sleep
class TimedFileHandler():
def __init__(self, file_name, method):
self._file_name = file_name
self._method = method
self._start = 0.0
def __enter__(self):
print(f"Entered file: {self._file_name}, in mode: {self._method}")
self._start = perf_counter()
self._file = open(self._file_name, self._method)
return self._file
def __exit__(self, exc_type, value, traceback):
if exception_type is not None:
print("Wrong filename or mode")
self._file.close()
end = perf_counter()
print("Exited file")
print(f"Time taken: {end - self._start}")
with TimedFileHandler('the-tales-of-beadle-the-bard.txt', 'w') as f:
f.write('C is simple.\n *(*(*var)())[7]X;')
sleep(5)
print(f"File closed?: {f.closed}")
Entered file: the-tales-of-beadle-the-bard.txt in mode: w Exited file Time taken: 5.001780649000011 File closed?: True
The same outcome can be achieved using by creating an object of the class:
handler = TimedFileHandler('the-tales-of-beadle-the-bard.txt', 'w')
f = handler.__enter__() # allocating
f.write('C is simple.\n *(*(*var)())[7]X;') # executing body
sleep(5)
handler.__exit__(None, None, None) # releasing
As a Generator
The contextmanager decorator from the contextlib module provides a more convenient way of implementing context managers. When the contextmanager decorator is used on a generator function, it automatically implements the required __enter__ and __exit__ methods and returns a context manager instead of the unusual iterator.
import contextlib
@contextlib.contextmanager
def timed_file_hanlder(file_name, method):
start = perf_counter()
file = open(file_name, method)
print(f"Entered file: {self._file_name}, in mode: {self._method}")
yield file
file.close()
end = perf_counter()
print("Exited file")
print(f"Time taken: {end - start}")
with timed_file_hanlder("the-tales-of-beadle-the-bard.txt","w") as f:
f.write("*Code written last month*\ I have no memory of this place.")
sleep(5)
print(f"File closed?: {f.closed}")
Entered file: the-tales-of-beadle-the-bard.txt in mode: w
Exited file
Time taken: 5.005866797999943
File closed?: True
Context Managers In Action
Django provides many useful context managers, such as transaction.atomic that enables developers to guarantee the atomicity of a database within a block of code. If the code in the with block raises an exception, all the changes made by the transaction are reverted.
from django.db import transaction with transaction.atomic(): # This code executes inside a transaction. do_something()
The Session class of the request module implements the __enter__ and __exit__ methods and can be used as a context manager when you need to preserve cookies between requests, want to make multiple requests to the same host or just want to keep TCP connection alive.
import requests with requests.Session() as sess: sess.request(method=method, url=url)
If you’ve read our previous articles on multithreading and multiprocessing, you would remember that we used the ThreadPoolExecutor and ProcessPoolExecutor context manager to automatically create and run threads/processes.
with concurrent.futures.ProcessPoolExecutor() as executor:
secs = [5, 4, 3, 2, 1]
pool = [executor.submit(useless_function, i) for i in secs]
for i in concurrent.futures.as_completed(pool):
print(f'Return Value: {i.result()}')
end = time.perf_counter()
print(f'Finished in {round(end-start, 2)} second(s)')
You can find the above code in a Colab notebook here.
Last Epoch
Python’s with statement provides a more convenient syntax when your code has to open and close connections or manage limited resources. There are more benefits to using the with statement than just a more concise syntax. The first advantage is that the whole allocation-release process happens under the control of the context manager object, making error prevention easier. Another reason to use it is that the with block makes it easier for developers to discern where the connection/resource is in use or can be used. Furthermore, it acts as a refactoring tool that moves the code for setup and teardown of any pair of operations to one place – the method definitions.
Context managers are one of those inherently Pythonic features that aren’t available in most languages. They provide a simple and elegant solution to a plethora of problems, even beyond resource management. I hope this article helped you better understand context managers so you can start using them to write better code.
If you want to learn more about context manager and the contextlib module, refer to the following resources:
