Python and the Context Managers

Context managers is one of those features of Python that makes it extremely easy to write clean code. Context managers are extensively used by Python programmers, but very few understand how they work. Let us explore them a bit.

Every python developer knows that the right way of working with files in Python is.

with open('myfile.txt', 'w') as my_file:
    do_something_with(my_file)

And of course we also know that this is equivalent to

my_file = open('myfile.txt', 'w')
try:
    do_something_with(my_file)
finally:
    my_file.close()

That is ok. But why and how does it work?

What is a context manager?

A context manager in Python is an object that implements __enter__(self) and __exit__(self, type, value, traceback) methods. Once our class implements those two methods, we can use it as a context manager with the with statement. But how does that work? When the control enters the with statement block, with statement will run the __enter__ method of the context manager. And once the code within the with block is done executing, the __exit__ method of the context manager is executed. We can demonstrate this by creating our own context manager. Let us define our class MyContextManager as follows.

class MyContextManager():
    def __enter__(self):
        print('The context has been setup.')
    def __exit__(self, type, value, traceback):
        print('Exiting the context.')

We use this class in a with statement as follows.

with MyContextManager():
    print('Within the context.')

The output of the statement is

The context has been setup.
Within the context.
Exiting the context.

This is how with statement makes sure that the resources are released once the block is complete. Let us see how we can make a context manager to work with files (ignoring the fact that the file descriptor already is a context manager).

Points to note before creating a context manager.

1. The __enter__ method of the context manager does not take any arguments. The value returned by the __enter__ method can be used with as statement within with statement (as in with open('myfile.txt', 'w') as f:).
2. The __exit__ method will take 3 arguments, the class type of the exception, the value of the exception and the traceback. All these arguments will be None if there is no exception raised in the with block.
3. If the __exit__ method return True, the with statement will suppress the error (does not propagate the error). Otherwise, the exception is thrown by the with block.

Now our context manager for file descriptor will be

class File():
    def __init__(self, name, mode):
        self.file = open(name, mode)

    def __enter__(self):
        print('The file will be returned.')
        return self.file

    def __exit__(self, type, value, traceback):
        print('The file will be closed now.')
        self.file.close()
        return True

This can be used with with statement like this.

with File('testfile.txt', 'w') as myfile:
    myfile.write('Hello world!')

The output on the console will is

The file will be returned.
The file will be closed now.

The result of the above statements is that a new file named testfile.txt will be opened (it will be created first if it does not exist already), the line Hello world! will be written to it. Once the with block is complete, the file will be closed.

Let us also examine how exceptions are handled by the context managers and with statement. Lets run the following statements.

with File('testfile.txt', 'w') as myfile:
    myfile.write('Hello world!')
    raise Exception('Could this exception create a problem?')

print('Exception is handled well')

Here we are raising an exception from within the with block (which is a possibility that we actually want to solve with context managers). The output is

The file will be returned.
The file will be closed now.
Exception is handled well

Since the __exit__ method returns True, the exception is not propagated by with. Hence the final print statement is executed. If the context manager is changed to

class File():
     def __init__(self, name, mode):
         self.file = open(name, mode)

     def __enter__(self):
         print('The file will be returned.')
         return self.file

     def __exit__(self, type, value, traceback):
         print('The file will be closed now.')
         self.file.close()
         return False

Notice that the __exit__ method returns False which instructs the with statement not to handle the exception (but propagate the exception). The result is

The file will be returned.
The file will be closed now.
---------------------------------------------------------------------------
Exception                                 Traceback (most recent call last)
<ipython-input-14-f7f6fb4fb916> in <module>()
      1 with File('testfile.txt', 'w') as myfile:
      2     myfile.write('Hello world!')
----> 3     raise Exception('Could this exception create a problem?')
      4 print('Exception is handled well')

Exception: Could this exception create a problem?

When the __exit__ returns False, the exception is propagated to the Python interpreter.

Creating context managers with contextlib

Python has a complete module called contextlib dedicated to create and manage context managers. We have already seen a way of creating context manager with a class implementing the __enter__ and __exit__ methods. Using the contextlib.contextmanager is another way to create a context manager.

The contextlib.contextmanager decorator can be used on an iterator function that yields exactly on value. It means the function must use exactly one yield statement. Every statement before the yield statement in the function is considered as part of the __enter__ method and everything after that is considered as part of the __exit__ method. The value that should be returned by the __enter__ method must be yielded. This is how we implement our File (our example class) context manager with the contextlib.contextmanager decorator.

from contextlib import contextmanager

@contextmanager
def file(name, mode):
    my_file = open(name, mode)
    print('The file is opened.')
    yield my_file
    print('The file will be closed now.')
    my_file.close()

Notice that we cannot return a value in a generator. So this is always equivalent to returning False in the __exit__ method. Now with our context manager file defined, let us run the following statements.

with file('testfile.txt', 'w') as myfile:
    myfile.write('Hellooooo world!')

The output is

The file is opened.
The file will be closed now.

The takeaway here is that the files opened are closed successfully. But how does it handle the exceptions raised in the with block. Running the statments

with file('testfile.txt', 'w') as myfile:
    myfile.write('Hellooooo world!')
    raise Exception('Could this cause a problem?')

print('Exception is handled well')

results in the following exception.

The file is opened.
---------------------------------------------------------------------------
Exception                                 Traceback (most recent call last)
<ipython-input-50-58f247e24470> in <module>()
      1 with file('testfile.txt', 'w') as myfile:
      2     myfile.write('Hellooooo world!')
----> 3     raise Exception('Could this cause a problem?')
      4 print('Exception is handled well')

Exception: Could this cause a problem?

Observe that the print statement post the yield statement is not executed. This exception will fail our context manager to close the file descriptors. A more sophisticated way to handle this is to rewrite the context manager as

from contextlib import contextmanager

@contextmanager                                        
def file(name, mode):             
    try:                      
        my_file = open(name, mode)
        print('The file is opened.')
        yield my_file                    
    finally:       
        print('The file will be closed now.')
        my_file.close()

Now the result of the previous statements is

The file is opened.
The file will be closed now.
---------------------------------------------------------------------------
Exception                                 Traceback (most recent call last)
<ipython-input-52-58f247e24470> in <module>()
    1 with file('testfile.txt', 'w') as myfile:
    2     myfile.write('Hellooooo world!')
----> 3     raise Exception('Could this cause a problem?')
    4 print('Exception is handled well')

Exception: Could this cause a problem?

This will successfully close the files even if an exception is raised. I hope this post helps you understand and use the context managers better. Cheers!