Mellowhost Blog

Context Manager in Python

In one line, context managers are an efficient way of handling resources in Python. So, what kind of resources are they? It could be any logical resource that you are using for your software, a common one, is database connections, or the files or in few cases, locks for concurrency control.

How exactly Context Manager is efficient?

If we talk about efficient programs, there could be several meanings of it in computer science. For our case, we mean efficient by writing less code, or more specifically, not writing repetitive codes for managing resources. There are some pinpoint benefits of using less repetitive codes other than writing more codes, which is purely technical. You do not forget to perform a step, that is essential if you do it from one source. Let’s find an example. You have a code, that connects to a remote FTP and uploads some content, once done, it closes the connection. If you do the process in several places, you might miss out to close the FTP connection in a place, that is accessed several times by the users. If that is so you might run out of the FTP connection pool on a random day. It is essential to close the connection after you are done with the connection to free the resources. Context Managers can help you write a code, that does the job for you, without the need to remember closing the connection each time.

The ‘with’ statement in Python

Before we can go deeper with Context Managers, we need to learn something about ‘with’ statement in python. ‘with’ is a special statement in python, that does jobs automatically for you. One, it calls a method of setting class, that is called ‘__enter__’ when it calls it, and the other, it automatically calls ‘__exit__’ method when it completes running its code. Let’s do some coding now

with open('test.txt') as f:
    f.write("testing")

In the above code, we are opening a file using the ‘with’ statement of python, then writing some texts in it. But wait, we haven’t closed the file, did you notice? Isn’t it necessary to close a file after opening it to free up the system file descriptors? Absolutely, but using the ‘with’ statement, do it for you even if you don’t do it in your code. In Python, the open() method for opening a file, can be used as Context Manager. For Context Manager, two methods are essential, one is when you set up the call, which is ‘__enter__’ and the other is, when you end the code, that is ‘__exit__’. ‘with’ statement in python is created to be used for Context Managers. As I said, the open() method can be used as Context Manager, which means, Python has both of these methods defined by default with the open() method, and can be used using the ‘with’ statement.

Why and when do we need Context Managers?

Before jumping into, how can we do context managers, let’s understand, if we understand the need of context managers properly with an example! The primary purpose of a context manager is to write cleaner nonrepetitive codes. Do we need this often? Yes, we do. A common example would be in setting up your database connections. If you are setting up a database connection and clean up the things once done, you may create a context manager to do that. There are more complex database use cases of context managers. Let’s focus on one of them. Let’s say, you would like to utilize the database ‘SAVEPOINT’ in a cleaner way, manage the release, and rollback for concurrent transactions based on the savepoint you create dynamically, a clean technique would be to use context manager. Pseudocode could be like the following for this kind of context manager

FUNCTION SAVEPOINT()
    NAME = UUID()
    SQL.EXECUTE('SAVEPOINT ' + UUID)
    TRY
        YIELD
    EXCEPTION
        EXECUTE('ROLLBACK TO SAVEPOINT ' + UUID)
        RAISE
    FINALLY
        EXECUTE('RELEASE SAVEPOINT ' + UUID)

What this context manager is trying to do, is generating a savepoint with a name for you. Once done, it yields the code you instruct it to run after the ‘with’ statement. If you create an exception from those codes, it rollback and sends the exception to the main program else if not it releases the savepoint and gives control to the code after the ‘with’ statement. This is technically the most efficient way of using Savepoint for SQL in your code. Similarly so, we can acquire and release locks for concurrent control using Context managers, or processing an API that had a setup call and an end/cleanup call.

How can we write Context Manager in Python?

There are two ways you can do it. One is using Python class, and the other use, using contextlib and the contextmanager decorator. Let’s first check out, how to do it using Python class to understand the concept better.

First, we want to emulate the way Python uses ‘open’ method as context manager using our own context manager class. A context manager class that can be used using ‘with’ statement could be like the following:

class Open_A_File():
    def __init__(self, name):
        self.name = name

    def __enter__(self):
        self.f = open(self.name)
        return self.f
    
    def __exit__(self, exc_type, exc_val, traceback):
        self.f.close()

with Open_A_File('test.txt') as f:
    f.write('Class Test')

In our class ‘Open_A_File()’, we have 3 methods. Our constructor __init__ method and other two methods are __enter__ and __exit__. When we used Open_A_File() using ‘with’ statement with a parameter, it setup our filename variable for the class using constructor, and then, calls the __enter__ method. It then opens the file and returns the file object. When it returns the file object, we catch it as ‘f’ to use in our code under the ‘with’ statement. We then write the code and the codes within ‘with’ statement ends, hence the __exit__ is automatically called, that closes the file object by calling ‘close()’. We can technically convert any class into a Context Manager and use them using ‘with’ statement if can define the methods to do while entering and exiting the class when called directly with the ‘with’ statement.

Other than class, we can use Context Managers using a function, through the use of contextlib library. This is the most used method of using Context Managers. We used the idea of this, in our pseudocode while demonstrating earlier. Let’s rewrite the above code using contextlib below

from contextlib import contextmanager

@contextmanager
def open_a_fiie(name):
    try:
        f = open(file)
        yield f
    finally:
        f.close()

with Open_A_File('test.txt') as f:
    f.write('Contextlib Test')

We first import contextmanager decorator from contextlib and then, we define a normal function. Although, there exists a ‘yield’ statement. For context managers, the statements that exist before yield would execute on __enter__ method, and the statements after yield would execute on __exit__ method. If you want to return anything to the ‘with’ call, then you need to specify that after yield, as we did in yield f, that means, we returned the file object to the ‘wite’ statement. the yield would replace the code, we run after the ‘with’ statement, like the f.write() in our case.

Hope this make sense. For confusion, or in case you would like to add some, do comment below. Thanks for reading.