Context Manager in Python

Managing Resources: In any programming language, the usage of resources like file operations or database connections is very common. But these resources are limited in supply. Therefore, the main problem lies in making sure to release these resources after usage. If they are not released then it will lead to resource leakage and may cause the system to either slow down or crash. It would be very helpful if users have a mechanism for the automatic setup and teardown of resources. In Python, it can be achieved by the usage of context managers which facilitate the proper handling of resources. The most common way of performing file operations is by using the keyword as shown below: 

Python3

# Python program showing
# a use of with keyword

with open("test.txt") as f:
    data = f.read()

Let’s take the example of file management. When a file is opened, a file descriptor is consumed which is a limited resource. Only a certain number of files can be opened by a process at a time. The following program demonstrates it. 

Python3

file_descriptors = []
for x in range(100000):
    file_descriptors.append(open('test.txt', 'w'))

Output:

Traceback (most recent call last):
  File "context.py", line 3, in 
OSError: [Errno 24] Too many open files: 'test.txt'

An error message saying that too many files are open. The above example is a case of file descriptor leakage. It happens because there are too many open files and they are not closed. There might be chances where a programmer may forget to close an opened file. 

Managing Resources using context manager: Suppose a block of code raises an exception or if it has a complex algorithm with multiple return paths, it becomes cumbersome to close a file in all the places. Generally in other languages when working with files try-except-finally is used to ensure that the file resource is closed after usage even if there is an exception. Python provides an easy way to manage resources: Context Managers. The with keyword is used. When it gets evaluated it should result in an object that performs context management. Context managers can be written using classes or functions(with decorators).

Creating a Context Manager: When creating context managers using classes, user need to ensure that the class has the methods: __enter__() and __exit__(). The __enter__() returns the resource that needs to be managed and the __exit__() does not return anything but performs the cleanup operations. First, let us create a simple class called ContextManager to understand the basic structure of creating context managers using classes, as shown below: 

Python3

# Python program creating a
# context manager

class ContextManager():
    def __init__(self):
        print('init method called')

    def __enter__(self):
        print('enter method called')
        return self

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

with ContextManager() as manager:
    print('with statement block')

Output:

init method called
enter method called
with statement block
exit method called

In this case, a ContextManager object is created. This is assigned to the variable after the keyword i.e manager. On running the above program, the following get executed in sequence:

• __init__()
• __enter__()
• statement body (code inside the with block)
• __exit__()[the parameters in this method are used to manage exceptions]

File management using context manager: Let’s apply the above concept to create a class that helps in file resource management. The FileManager class helps in opening a file, writing/reading contents, and then closing it. 

Python3

# Python program showing
# file management using
# context manager

class FileManager():
    def __init__(self, filename, mode):
        self.filename = filename
        self.mode = mode
        self.file = None

    def __enter__(self):
        self.file = open(self.filename, self.mode)
        return self.file

    def __exit__(self, exc_type, exc_value, exc_traceback):
        self.file.close()

# loading a file
with FileManager('test.txt', 'w') as f:
    f.write('Test')

print(f.closed)

Output:

True

File management using context manager and with statement: On executing the with block, the following operations happen in sequence:

• A FileManager object is created with test.txt as the filename and w(write) as the mode when __init__ method is executed.
• The __enter__ method opens the test.txt file in write mode(setup operation) and returns a file object to variable f.
• The text ‘Test’ is written into the file.
• The __exit__ method takes care of closing the file on exiting the with block(teardown operation). When print(f.closed) is run, the output is True as the FileManager has already taken care of closing the file which otherwise needed to be explicitly done.

Database connection management using context manager: Let’s create a simple database connection management system. The number of database connections that can be opened at a time is also limited(just like file descriptors). Therefore context managers are helpful in managing connections to the database as there could be chances that the programmer may forget to close the connection. 

Python3

# Python program shows the
# connection management
# for MongoDB

from pymongo import MongoClient

class MongoDBConnectionManager():
    def __init__(self, hostname, port):
        self.hostname = hostname
        self.port = port
        self.connection = None

    def __enter__(self):
        self.connection = MongoClient(self.hostname, self.port)
        return self.connection

    def __exit__(self, exc_type, exc_value, exc_traceback):
        self.connection.close()

# connecting with a localhost
with MongoDBConnectionManager('localhost', '27017') as mongo:
    collection = mongo.connection.SampleDb.test
    data = collection.find({'_id': 1})
    print(data.get('name'))

Database connection management using context manager and with statement: On executing the with block, the following operations happen in sequence:

• A MongoDBConnectionManager object is created with localhost as the hostname name and 27017 as the port when the __init__ method is executed.
• The __enter__ method opens the MongoDB connection and returns the MongoClient object to variable mongo.
• The test collection in the SampleDb database is accessed and the document with _id=1 is retrieved. The name field of the document is printed.
• The __exit__ method takes care of closing the connection on exiting the with block(teardown operation).