Context Manager in Python

The usage of resources such as files, locks, database connections etc. is very common in a program. Since these resources are limited in supply, we should always make sure to release them after usage. If not, then the program will keep the resources alive thereby compromising valuable system resources, such as memory (memory leak) and network bandwidth, which in turn causes the system to either slow down or crash.

Context Manager is a mechanism in Python that facilitates the proper handling of resources through automatic setup and teardown of resources.

Consider an example program for writing some content to a file. Here, the file is our resource. Let's look at this process in 3 phases:

1. Setup phase - We acquire the resource in this phase. As per our example, we will open the file in this phase.

2. Processing phase - We use the acquired resource in this phase. As per our example, we will write the content to the opened file.

3. Teardown phase - This is the final phase where we do the cleanup actions to release the resource. As per our example, we will close the file in this phase.

Let's translate this into code 💻:

# Setup
file = open('path/to/my_file.txt', 'w')
# Process
file.write('Hello Python')
# Teardown
file.close()

Here we are releasing the resource after use, which is good, but there is one problem though. What if some Exception was thrown from the processing phase 😕? In that case, the logic written for teardown won't be executed and hence the resource won't be released. How can we fix this? 🤔

Try ... Finally

We know that the finally block always executes, regardless of whether an exception is thrown. So we can move the teardown actions into the finally block 💡:

file = open('path/to/my_file.txt', 'w')
try:
    file.write('Hello Python')
except Exception as error:
    print(f'Error: {error} occured while writing to the file')
finally:
    file.close()

This is one way of doing it but this will become cumbersome if we have to manage resources at multiple places in our program. Context Manager provides an easy way to manage resources using with statement.

With Statement

The with statement creates a runtime context in which the resource is available, and then automatically releases the resource when the block of code is finished. The general syntax is as follows 💡:

with expression as target_var:
    process(target_var)

The context manager object results from evaluating the expression after with. In other words, expression must return an object that implements the context management protocol. This protocol consists of two special methods:

1. __enter__() is called by the with statement to enter the runtime context. (setup)

2. __exit__() is called when the execution leaves the with block. (teardown)

The as specifier is optional. If we provide a target_var with as, then the return value of calling __enter__() on the context manager object is bound to that variable.

With this knowledge, let's refactor our file management program to use context manager instead of try-finally 💻:

with open('path/to/my_file.txt', 'w') as file:
    file.write('Hello Python')

When we run this with statement, open() returns an io.TextIOBase object. Since this object is also a context manager, the with statement calls __enter__() and assigns its return value to file. Then we can write content to the file inside the block by using this file object. When the block ends, __exit__() is automatically called which in turn closes the file, even if an exception is raised inside the with block.

Compared to try...finally constructs, this can make our code much clearer, safer, and reusable and hence you will be seeing more of this syntax in the libraries/code bases that you will be working on.

Creating Context Manager using Class

Creating a context manager using class is very easy. All we have to do is define the special methods: __enter__() and __exit__(). ✌️

If the __exit__() returns True, then any exception that occurs in the with block is swallowed and the execution continues at the next statement after with. Otherwise, exceptions are propagated out of the context. We can take advantage of this feature to encapsulate exception handling inside the context manager.

First, let us create a simple ContextManager class to understand the basic structure of creating context managers using classes 😃:

class ContextManager:
    def __init__(self):
        print('Initialized')

    def __enter__(self):
        print('Setup')
        return self

    def __exit__(self, exc_type, exc_value, exc_traceback):
        print('Teardown')

with ContextManager() as manager:
    print('Processed')

# Initialized
# Setup
# Processed
# Teardown

Now that we know how to create a context manager using class and its working, let's practice creating one to manage database connection 🥸:

class DatabaseManager:
    def __init__(self, host='localhost', port=3306):
        self.host = host
        self.port = port
        self.connection = None

    def __enter__(self):
        # self.connection = connect(self.host, self.port)
        self.connection = {'query': lambda q: [1, 2, 3]}
        print(f'Connected to {self.host}:{self.port}')
        return self.connection

    def __exit__(self, exc_type, exc_value, exc_traceback):
        # self.connection.close()
        print("Connection closed")

with DatabaseManager('my_server.net') as db:
    results = db['query']('SELECT id FROM my_table')
    print(f'Ids: {results}')

# Connected to my_server.net:3306
# Ids: [1, 2, 3]
# Connection closed

Creating Context Manager using Function

Python’s generator functions and the contextlib.contextmanager decorator provides an alternative way to implement the context management protocol. If we decorate an appropriately coded generator function with @contextmanager, then we get a function-based context manager that automatically provides both required methods: __enter__() and __exit__(). This is how our DatabaseManager would look if implemented using the function-based approach 💻:

from contextlib import contextmanager

@contextmanager
def DatabaseManager(host='localhost', port=3306):
     print(f'Connected to {host}:{port}')
     yield {'query': lambda q: [1, 2, 3]}
     print("Connection closed")

with DatabaseManager('my_server.net') as db:
    results = db['query']('SELECT id FROM my_table')
    print(f'Ids: {results}')

# Connected to my_server.net:3306
# Ids: [1, 2, 3]
# Connection closed

We will look at the working of this, after exploring generator functions in Python. 👍

---

Well, that’s all for this article. Thanks for reading! 🙏

I hope you liked this article and if you have any questions or suggestions, please drop them in the comments. Happy Coding.. 👋