The contextlib
module in Python provides a number of utilities to help in creating and managing context managers.
Context managers offers a more convenient way to automate setup and teardown operations.
Typically, context manger objects are used with the with
statements. They are marked by two methods i.e __enter__()
and __exit__()
.
The __enter__()
method is responsible for setting up and allocating the necessary resources for the object, while, the __exit__()
method is responsible for cleaning up and freeing up all the resources used by the object.
The following example shows a simple example of a context manager object.
class MyContext:
def __init__(self):
print('__int__() called')
def __str__(self):
return "MyObject()"
#define the __enter__ method
def __enter__(self):
print('__enter__() called')
#define the __exit__ method
def __exit__(self, *args, **kwargs):
print('__exit__() called')
with MyContext() as ctx:
print('Do Something')
As you can see above, when such an object is used with the with statement, the two methods __enter__()
and __exit__()
are called automatically. The __enter__()
method is called when the control flow enters the with
block while the __exit__()
method is called immediately before the control flow leaves the with block.
A common example of context manager objects are the file objects created using the builtin open() function. When we use the file objects with the with statement, the file is automatically closed after the block of code within the with
statement has been executed, this alleviates the need to call the file.close()
method manually.
#the file path
myfile = 'demo.txt'
#open the file using the with statement
with open(myfile) as file:
#do something with the file
for line in file.readlines():
print(line)
#the file is now closed
Transform a generator function to a context manager
Generator functions are one of the common cases where setup and teardown operations maybe necessary.
In situations where you have a generator function, it may be trivial to write a new context manager class in the traditional way by defining __enter__()
and __exit__()
methods. The @contextmanager
decorator in the contextlib
module allows you to transform a generator function into a context manager in a convenient way and without having to explicitly define the __enter__()
and __exit__()
methods.
from contextlib import contextmanager
@contextmanager
def square(n):
'''a generator function to get a list of squares of numbers from 0 to n'''
try:
data = []
for i in range(n):
data.append(i ** 2)
yield data
except RuntimeError as e:
print('An erro occured')
finally:
print('Done')
with square(5) as sq:
for s in sq:
print(s)
The generator function will have to yield
exactly once, the value yielded will be assigned to the variable specified by the as
clause. Any exception raised inside the with
block are sent to the generator to be handled.
Objects with a close() method
If an object needs to be closed after a task is done and has defined the close()
method, we can use the closing()
function to close the object automatically even if it does not have the __enter__()
and __exit__()
methods defined.
from contextlib import closing
class School:
def __init__(self):
print('__init__() called')
def close(self):
print('close() called')
with closing(School()) as sch:
print('inside with block')